add PBD id to context

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java

@@ -61,6 +61,9 @@ interface Context {
     /** A client for handling state requests. */
     BeamFnStateClient getBeamFnStateClient();
 
+    /** The id of the parent ProcessBundleDescriptor. */
+    String getProcessBundleDescriptorId();
+
     /** The id of the PTransform. */
     String getPTransformId();
 

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java

@@ -319,6 +319,11 @@ public BeamFnStateClient getBeamFnStateClient() {
                       return beamFnStateClient;
                     }
 
+                    @Override
+                    public String getProcessBundleDescriptorId() {
+                      return processBundleDescriptor.getId();
+                    }
+
                     @Override
                     public String getPTransformId() {
                       return pTransformId;

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/debug/DataSampler.java

@@ -28,8 +28,10 @@
 import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 public class DataSampler {
   public static class OutputSampler<T> {
@@ -92,24 +94,45 @@ private void clear() {
     }
   }
 
-  private final Map<String, OutputSampler<?>> outputSamplers = new HashMap<>();
+  private final Map<String, Map<String, OutputSampler<?>>> outputSamplers = new HashMap<>();
 
-  public <T>OutputSampler<T> sampleOutput(String outputName, Coder<T> coder) {
-    if (outputSamplers.containsKey(outputName)) {
-      return (OutputSampler<T>)outputSamplers.get(outputName);
-    }
+  public <T>OutputSampler<T> sampleOutput(String processBundleDescriptorId, String pcollectionId, Coder<T> coder) {
+    outputSamplers.putIfAbsent(processBundleDescriptorId, new HashMap<>());
+    Map<String, OutputSampler<?>> samplers = outputSamplers.get(processBundleDescriptorId);
+    samplers.putIfAbsent(pcollectionId, new OutputSampler<T>(coder));
 
-    OutputSampler<T> sampler = new OutputSampler<T>(coder);
-    outputSamplers.put(outputName, sampler);
-    return sampler;
+    return (OutputSampler<T>)samplers.get(pcollectionId);
   }
 
   public Map<String, List<byte[]>> samples() {
+    return samplesFor(new HashSet<>(), new HashSet<>());
+  }
+
+  public Map<String, List<byte[]>> samplesFor(Set<String> descriptors, Set<String> pcollections) {
     Map<String, List<byte[]>> samples = new HashMap<>();
-    outputSamplers.forEach((pcollectionId, outputSampler) -> {
-      samples.putIfAbsent(pcollectionId, new ArrayList<>());
-      samples.get(pcollectionId).addAll(outputSampler.samples());
+    outputSamplers.forEach((descriptorId, samplers) -> {
+      if (!descriptors.isEmpty() && !descriptors.contains(descriptorId)) {
+        return;
+      }
+
+      samplers.forEach((pcollectionId, outputSampler) -> {
+        if (!pcollections.isEmpty() && !pcollections.contains(pcollectionId)) {
+          return;
+        }
+
+        samples.putIfAbsent(pcollectionId, new ArrayList<>());
+        samples.get(pcollectionId).addAll(outputSampler.samples());
+      });
     });
+
     return samples;
   }
+
+  public Map<String, List<byte[]>> samplesForDescriptors(Set<String> descriptors) {
+    return samplesFor(descriptors, new HashSet<>());
+  }
+
+  public Map<String, List<byte[]>> samplesForPCollections(Set<String> pcollections) {
+    return samplesFor(new HashSet<>(), pcollections);
+  }
 }

sdks/java/harness/src/test/java/org/apache/beam/fn/harness/debug/DataSamplerTest.java

@@ -7,6 +7,7 @@
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarIntCoder;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -23,71 +24,13 @@
 
 @RunWith(JUnit4.class)
 public class DataSamplerTest {
-
     @Test
     public void testCreatingAndProcessingWithSampling() throws Exception {
+        DataSampler sampler = new DataSampler();
 
+        VarIntCoder coder = VarIntCoder.of();
+        sampler.sampleOutput("descriptor-id", "pcollection-id", coder).sample(12345);
 
-        // Create the DataSampling PTransform.
-        String pTransformId = "pTransformId";
-
-        RunnerApi.FunctionSpec functionSpec =
-                RunnerApi.FunctionSpec.newBuilder()
-                        .setUrn(DataSamplingFnRunner.URN)
-                        .build();
-        RunnerApi.PTransform pTransform =
-                RunnerApi.PTransform.newBuilder()
-                        .setSpec(functionSpec)
-                        .putInputs("input", "inputTarget")
-                        .build();
-
-        // Populate fake input PCollections.
-        Map<String, RunnerApi.PCollection> pCollectionMap = new HashMap<>();
-        pCollectionMap.put(
-                "inputTarget",
-                RunnerApi.PCollection.newBuilder()
-                        .setUniqueName("inputTarget")
-                        .setCoderId("coder-id")
-                        .build());
-
-        // Populate the PTransform context that includes the DataSampler.
-        DataSampler dataSampler = new DataSampler();
-        RunnerApi.Coder coder = CoderTranslation.toProto(StringUtf8Coder.of()).getCoder();
-        PTransformRunnerFactoryTestContext context =
-                PTransformRunnerFactoryTestContext.builder(pTransformId, pTransform)
-                        .processBundleInstructionId("instruction-id")
-                        .pCollections(pCollectionMap)
-                        .coders(Collections.singletonMap("coder-id", coder))
-                        .dataSampler(dataSampler)
-                        .build();
-
-        // Create the runner which samples the input PCollection.
-        new DataSamplingFnRunner.Factory<>().createRunnerForPTransform(context);
-        assertThat(
-                context.getPCollectionConsumers().keySet(),
-                contains("inputTarget"));
-
-        // Send in a test value that should be sampled.
-        context.getPCollectionConsumer("inputTarget").accept(valueInGlobalWindow("Hello, World!"));
-
-        // Rehydrate the given utf-8 string coder.
-        RehydratedComponents rehydratedComponents =
-                RehydratedComponents.forComponents(
-                                RunnerApi.Components.newBuilder()
-                                        .putAllCoders(context.getCoders())
-                                        .putAllPcollections(context.getPCollections())
-                                        .putAllWindowingStrategies(context.getWindowingStrategies())
-                                        .build())
-                        .withPipeline(Pipeline.create());
-        Coder<String> rehydratedCoder = (Coder<String>)rehydratedComponents.getCoder("coder-id");
-
-        Map<String, List<byte[]>> samples = dataSampler.samples();
-        assertThat(samples.keySet(), contains("inputTarget"));
-
-        // Ensure that the value was sampled.
-        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
-        rehydratedCoder.encode("Hello, World!", outputStream);
-        byte[] encodedValue = outputStream.toByteArray();
-        assertThat(samples.get("inputTarget"), contains(encodedValue));
+        System.out.println(sampler.samples());
     }
 }