vishalmore90 opened a new issue, #41038: URL: https://github.com/apache/beam/issues/41038
### What happened? In `org.apache.beam.runners.spark.translation.ValueAndCoderLazySerializable`, `writeCommon(OutputStream out)` attempts to optimize serialization buffer allocation by checking if the element's coder can cheaply observe byte sizes via `coder.isRegisterByteSizeObserverCheap(value)`. If true, it attempts to pre-size the `ByteArrayOutputStream` using a custom `ByteSizeObserver`: https://github.com/apache/beam/blob/master/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/ValueAndCoderLazySerializable.java#L86-L115 ```java private static class ByteSizeObserver extends ElementByteSizeObserver { private long observedSize = 0; @Override protected void reportElementSize(long elementByteSize) { observedSize += elementByteSize; } } void writeCommon(OutputStream out) throws IOException { if (!(coderOrBytes instanceof Coder)) { byte[] bytes = (byte[]) coderOrBytes; VarInt.encode(bytes.length, out); out.write(bytes); } else { @SuppressWarnings("unchecked") Coder<T> coder = (Coder<T>) coderOrBytes; int bufferSize = 1024; if (coder.isRegisterByteSizeObserverCheap(value)) { try { ByteSizeObserver observer = new ByteSizeObserver(); coder.registerByteSizeObserver(value, observer); bufferSize = (int) observer.observedSize; } catch (Exception e) { Throwables.throwIfUnchecked(e); throw new RuntimeException(e); } } ByteArrayOutputStream bytes = new ByteArrayOutputStream(bufferSize); try { coder.encode(value, bytes); ... ``` However, `ByteSizeObserver.reportElementSize(long)` is **never invoked** because `observer.advance()` is never called. #### Root Cause `ElementByteSizeObserver` works in two phases: 1. `coder.registerByteSizeObserver(value, observer)` notifies the observer via `observer.update(size)`. In `ElementByteSizeObserver.java` (lines 67–75), `update()` only accumulates into the private field `totalSize`: ```java @Override public void update(Observable obs, Object obj) { if (obj instanceof Long) { totalSize += (long) (scalingFactor * (Long) obj); } else if (obj instanceof Integer) { totalSize += (long) (scalingFactor * (Integer) obj); } ... } ``` 2. The abstract method `reportElementSize(long)` is only called when `observer.advance()` is invoked (`ElementByteSizeObserver.java` lines 81–85): ```java public void advance() { reportElementSize(totalSize); totalSize = 0; isLazy = false; } ``` Because `writeCommon` never calls `observer.advance()`, `ByteSizeObserver.reportElementSize` is never executed. Consequently: - `observer.observedSize` is always `0`. - `bufferSize = (int) observer.observedSize;` assigns `bufferSize = 0`. - `ByteArrayOutputStream bytes = new ByteArrayOutputStream(0);` is allocated with an internal buffer of capacity 0 (`new byte[0]`). ### Code Pointers / Steps to Reproduce 1. Compare with all other implementations and usages of `ElementByteSizeObserver` across Apache Beam, which all properly invoke `advance()`: - `sdks/java/core/.../GroupIntoBatches.java#L151-L154`: ```java ByteSizeObserver observer = new ByteSizeObserver(); valueCoder.registerByteSizeObserver(element, observer); observer.advance(); return observer.getElementByteSize(); ``` - `sdks/java/harness/.../PCollectionConsumerRegistry.java#L533-L538`: ```java byteCountObserver = new ByteSizeObserver(); coder.registerByteSizeObserver(value, byteCountObserver); if (!byteCountObserver.getIsLazy()) { byteCountObserver.advance(); this.distribution.update(byteCountObserver.observedSize); } ``` - `sdks/java/io/google-cloud-platform/.../bigtable/.../CoderSizeEstimator.java#L63-L64` and `spanner/.../SizeEstimator.java#L60-L61`: ```java coder.registerByteSizeObserver(element, observer); observer.advance(); ``` - `runners/google-cloud-dataflow-java/.../PartialGroupByKeyParDoFns.java#L276-L281`: ```java Observer observer = new Observer(); coder.registerByteSizeObserver(value, observer); if (!observer.getIsLazy()) { observer.advance(); return observer.observedSize; } ``` 2. Minimal reproduction: ```java BigEndianIntegerCoder coder = BigEndianIntegerCoder.of(); ValueAndCoderLazySerializable<Integer> item = ValueAndCoderLazySerializable.of(42, coder); // In writeCommon(): // coder.isRegisterByteSizeObserverCheap(42) is true. // coder.registerByteSizeObserver(42, observer) calls observer.update(4L). // Without observer.advance(), observer.observedSize remains 0. // bufferSize becomes 0. ``` 3. Existing testing gap: In `runners/spark/src/test/java/org/apache/beam/runners/spark/translation/ValueAndCoderLazySerializableTest.java`, all unit tests use `IterableCoder`, for which `isRegisterByteSizeObserverCheap()` is hardcoded to `false` (`IterableLikeCoder.java#L181-L183`). The `isRegisterByteSizeObserverCheap` branch has zero test coverage. ### Impact This defect causes an anti-optimization on the critical shuffle and persistence serialization path of the Spark runner: - For coders where `isRegisterByteSizeObserverCheap` is `false`, `bufferSize` defaults to `1024` bytes. - For coders where `isRegisterByteSizeObserverCheap` is `true` (e.g. `BigEndianIntegerCoder`, `BigEndianLongCoder`, `VarIntCoder`, `VarLongCoder`, `ByteArrayCoder`, `StringUtf8Coder`, `KvCoder`, `NullableCoder`, etc.), `bufferSize` is set to **`0`**. - In Spark pipelines using `ValueAndCoderLazySerializable` (such as RDD persistence in `BoundedDataset.java:119`, accumulator serialization in `GroupCombineFunctions.java`, and multi-output routing in `TranslationUtils.java`), serializing every single cheap element allocates a 0-capacity `ByteArrayOutputStream(0)`. - The very first byte written to `ByteArrayOutputStream` forces an immediate buffer reallocation and array copy, followed by repeated power-of-two growth copies for larger byte arrays or strings. ### Proposed Solution 1. In `runners/spark/src/main/java/org/apache/beam/runners/spark/translation/ValueAndCoderLazySerializable.java`, check `!observer.getIsLazy()` and call `observer.advance()` before reading `observer.observedSize`. 2. Guard against non-positive sizes so that `bufferSize` retains a sensible minimum (such as the default 1024 bytes): ```java int bufferSize = 1024; if (coder.isRegisterByteSizeObserverCheap(value)) { try { ByteSizeObserver observer = new ByteSizeObserver(); coder.registerByteSizeObserver(value, observer); if (!observer.getIsLazy()) { observer.advance(); if (observer.observedSize > 0) { bufferSize = (int) observer.observedSize; } } } catch (Exception e) { Throwables.throwIfUnchecked(e); throw new RuntimeException(e); } } ByteArrayOutputStream bytes = new ByteArrayOutputStream(bufferSize); ``` 3. Add test coverage in `ValueAndCoderLazySerializableTest.java` verifying serialization with a cheap coder (e.g. `BigEndianIntegerCoder.of()` or `ByteArrayCoder.of()`). ### Issue Priority Priority: 2 (default / most bugs should be filed as P2) ### Issue Components - [ ] Component: Python SDK - [ ] Component: Java SDK - [ ] Component: Go SDK - [ ] Component: Typescript SDK - [ ] Component: IO connector - [ ] Component: Beam YAML - [ ] Component: Beam examples - [ ] Component: Beam playground - [ ] Component: Beam katas - [ ] Component: Website - [ ] Component: Infrastructure - [x] Component: Spark Runner - [ ] Component: Flink Runner - [ ] Component: Prism Runner - [ ] Component: Hazelcast Jet Runner - [ ] Component: Google Cloud Dataflow Runner -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
