Calculate byte size via sampling in StateBackedIterable if size is not cheap to calculate

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateBackedIterable.java

@@ -30,6 +30,8 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.function.Supplier;
 import org.apache.beam.fn.harness.Cache;
 import org.apache.beam.fn.harness.Caches;
@@ -106,6 +108,11 @@ private static class WrappedObservingIterator<T> extends ElementByteSizeObservab
     private boolean observerNeedsAdvance = false;
     private boolean exceptionLogged = false;
 
+    // Lowest sampling probability: 0.001%.
+    private static final int SAMPLING_TOKEN_UPPER_BOUND = 1000000;
+    private static final int SAMPLING_CUTOFF = 10;
+    private int samplingToken = 0;
+
     static <T> WrappedObservingIterator<T> create(
         Iterator<T> iterator, org.apache.beam.sdk.coders.Coder<T> elementCoder) {
       WrappedObservingIterator<T> result = new WrappedObservingIterator<>(iterator, elementCoder);
@@ -125,6 +132,18 @@ private WrappedObservingIterator(
       this.elementCoder = elementCoder;
     }
 
+    private boolean sampleElement() {
+      // Sampling probability decreases as the element count is increasing.
+      // We unconditionally sample the first samplingCutoff elements. For the
+      // next samplingCutoff elements, the sampling probability drops from 100%
+      // to 50%. The probability of sampling the Nth element is:
+      // min(1, samplingCutoff / N), with an additional lower bound of
+      // samplingCutoff / samplingTokenUpperBound. This algorithm may be refined
+      // later.
+      samplingToken = Math.min(samplingToken + 1, SAMPLING_TOKEN_UPPER_BOUND);
+      return ThreadLocalRandom.current().nextInt(samplingToken) < SAMPLING_CUTOFF;
+    }
+
     @Override
     public boolean hasNext() {
       if (observerNeedsAdvance) {
@@ -138,15 +157,20 @@ public boolean hasNext() {
     public T next() {
       T value = wrappedIterator.next();
       try {
-        elementCoder.registerByteSizeObserver(value, observerProxy);
-        if (observerProxy.getIsLazy()) {
-          // The observer will only be notified of bytes as the result
-          // is used. We defer advancing the observer until hasNext in an
-          // attempt to capture those bytes.
-          observerNeedsAdvance = true;
-        } else {
-          observerNeedsAdvance = false;
-          observerProxy.advance();
+        boolean cheap = elementCoder.isRegisterByteSizeObserverCheap(value);
+        if (cheap || sampleElement()) {
+          observerProxy.setScalingFactor(
+              cheap ? 1.0 : Math.max(samplingToken, SAMPLING_CUTOFF) / (double) SAMPLING_CUTOFF);
+          elementCoder.registerByteSizeObserver(value, observerProxy);
+          if (observerProxy.getIsLazy()) {
+            // The observer will only be notified of bytes as the result
+            // is used. We defer advancing the observer until hasNext in an
+            // attempt to capture those bytes.
+            observerNeedsAdvance = true;
+          } else {
+            observerNeedsAdvance = false;
+            observerProxy.advance();
+          }
         }
       } catch (Exception e) {
         if (!exceptionLogged) {

sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateBackedIterableTest.java

@@ -269,7 +269,9 @@ public void testByteObservingStateBackedIterable() throws Exception {
               .sum();
       observer.advance();
       // 5 comes from size and hasNext (see IterableLikeCoder)
-      assertEquals(iterateBytes + 5, observer.total);
+      // observer receives scaled, StringUtf8Coder is not cheap so sampling may produce value that
+      // is off
+      assertEquals((float) iterateBytes + 5, (float) observer.total, 3);
     }
   }