This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 023ca6d9a63 [Java] Cache DoFn invoker constructors per instance
(#39310)
023ca6d9a63 is described below
commit 023ca6d9a638750ea76e0b63cb6aa366d7bc6204
Author: Alexander Kolb <[email protected]>
AuthorDate: Mon Jul 20 19:42:22 2026 +0200
[Java] Cache DoFn invoker constructors per instance (#39310)
* Cache DoFn invoker constructors per instance
Co-Authored-By: Codex <[email protected]>
* Add DoFn invoker cache fix to changelog
Co-Authored-By: Codex <[email protected]>
* Clarify DoFn cache performance improvement
Co-Authored-By: Codex <[email protected]>
---------
Co-authored-by: Codex <[email protected]>
---
CHANGES.md | 1 +
.../reflect/ByteBuddyDoFnInvokerFactory.java | 83 ++++++++++++++--------
.../sdk/transforms/reflect/DoFnInvokersTest.java | 47 ++++++++++++
3 files changed, 103 insertions(+), 28 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 19a16e1dcae..d39a62f1a78 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -83,6 +83,7 @@
## Bugfixes
* Fixed unbounded checkpoint state growth for splittable DoFns that
self-checkpoint on the portable Flink runner (Java)
([#27648](https://github.com/apache/beam/issues/27648)).
+* Improved Java pipeline performance by avoiding repeated `DoFn` type
descriptor resolution when creating cached invokers
([#39309](https://github.com/apache/beam/issues/39309)).
* Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
## Security Fixes
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
index c08243fda5c..14f864c5b04 100644
---
a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
+++
b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
@@ -29,6 +29,7 @@ import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
import net.bytebuddy.ByteBuddy;
import net.bytebuddy.description.field.FieldDescription;
import net.bytebuddy.description.method.MethodDescription;
@@ -113,6 +114,7 @@ import org.apache.beam.sdk.util.UserCodeException;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.sdk.values.TypeDescriptors;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.MapMaker;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.Primitives;
import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
@@ -242,6 +244,22 @@ class ByteBuddyDoFnInvokerFactory implements
DoFnInvokerFactory {
private final Map<InvokerCacheKey, Constructor<?>>
byteBuddyInvokerConstructorCache =
new ConcurrentHashMap<>();
+ /**
+ * A weak, identity-keyed cache of the constructor selected for each {@link
DoFn} instance.
+ *
+ * <p>Selecting a constructor resolves the {@link DoFn}'s input and output
type descriptors. Some
+ * runners create a new invoker for the same {@link DoFn} at every bundle
boundary, so doing that
+ * work on every invocation can be expensive. A deserialized {@link DoFn} is
a new instance, so
+ * its constructor is selected again; later invoker creation for that
instance reuses the cached
+ * selection.
+ *
+ * <p>The constructor values do not reference their {@link DoFn} keys, and
weak keys allow unused
+ * instances to be collected. {@link MapMaker#weakKeys} also uses identity
equality, which is
+ * required because separate instances of the same class may have different
generic types.
+ */
+ private final ConcurrentMap<DoFn<?, ?>, Constructor<?>>
byteBuddyInvokerConstructorInstanceCache =
+ new MapMaker().weakKeys().makeMap();
+
private ByteBuddyDoFnInvokerFactory() {}
/** Returns the {@link DoFnInvoker} for the given {@link DoFn}. */
@@ -330,39 +348,15 @@ class ByteBuddyDoFnInvokerFactory implements
DoFnInvokerFactory {
signature.fnClass(),
fn.getClass());
- // Extract input and output type descriptors to distinguish generic
instantiations.
- // Fall back to Object.class if unavailable. When type info is lost,
different generic
- // instantiations share an invoker, which is acceptable since the DoFn
class in the cache
- // key prevents collisions between different DoFn classes.
- TypeDescriptor<InputT> inputType;
- try {
- inputType = fn.getInputTypeDescriptor();
- } catch (Exception e) {
- // Some DoFns (like MapElements) throw IllegalStateException if queried
after
- // serialization.
- // In this case, we fall back to the raw class behavior (Object).
- inputType = null;
- }
- if (inputType == null) {
- inputType = (TypeDescriptor<InputT>) TypeDescriptor.of(Object.class);
- }
-
- TypeDescriptor<OutputT> outputType;
- try {
- outputType = fn.getOutputTypeDescriptor();
- } catch (Exception e) {
- // Same as above: fall back to Object if type info is unavailable.
- outputType = null;
- }
- if (outputType == null) {
- outputType = (TypeDescriptor<OutputT>) TypeDescriptor.of(Object.class);
- }
+ Constructor<?> invokerConstructor =
+ byteBuddyInvokerConstructorInstanceCache.computeIfAbsent(
+ fn, unused -> getByteBuddyInvokerConstructor(signature, fn));
try {
@SuppressWarnings("unchecked")
DoFnInvokerBase<InputT, OutputT, DoFn<InputT, OutputT>> invoker =
(DoFnInvokerBase<InputT, OutputT, DoFn<InputT, OutputT>>)
- getByteBuddyInvokerConstructor(signature, inputType,
outputType).newInstance(fn);
+ invokerConstructor.newInstance(fn);
if (signature.onTimerMethods() != null) {
for (OnTimerMethod onTimerMethod :
signature.onTimerMethods().values()) {
@@ -389,6 +383,39 @@ class ByteBuddyDoFnInvokerFactory implements
DoFnInvokerFactory {
}
}
+ private <InputT, OutputT> Constructor<?> getByteBuddyInvokerConstructor(
+ DoFnSignature signature, DoFn<InputT, OutputT> fn) {
+ // Extract input and output type descriptors to distinguish generic
instantiations.
+ // Fall back to Object.class if unavailable. When type info is lost,
different generic
+ // instantiations share an invoker, which is acceptable since the DoFn
class in the cache
+ // key prevents collisions between different DoFn classes.
+ TypeDescriptor<InputT> inputType;
+ try {
+ inputType = fn.getInputTypeDescriptor();
+ } catch (Exception e) {
+ // Some DoFns (like MapElements) throw IllegalStateException if queried
after
+ // serialization.
+ // In this case, we fall back to the raw class behavior (Object).
+ inputType = null;
+ }
+ if (inputType == null) {
+ inputType = (TypeDescriptor<InputT>) TypeDescriptor.of(Object.class);
+ }
+
+ TypeDescriptor<OutputT> outputType;
+ try {
+ outputType = fn.getOutputTypeDescriptor();
+ } catch (Exception e) {
+ // Same as above: fall back to Object if type info is unavailable.
+ outputType = null;
+ }
+ if (outputType == null) {
+ outputType = (TypeDescriptor<OutputT>) TypeDescriptor.of(Object.class);
+ }
+
+ return getByteBuddyInvokerConstructor(signature, inputType, outputType);
+ }
+
/**
* Returns a generated constructor for a {@link DoFnInvoker} for the given
{@link DoFnSignature}
* and specific generic types.
diff --git
a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java
b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java
index e9269deb18e..9f4fdb6c9c2 100644
---
a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java
+++
b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java
@@ -82,6 +82,7 @@ import org.apache.beam.sdk.values.OutputBuilder;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.sdk.values.TypeDescriptors;
import org.apache.beam.sdk.values.WindowedValues;
+import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Instant;
import org.junit.Before;
import org.junit.Rule;
@@ -1440,6 +1441,18 @@ public class DoFnInvokersTest {
public TypeDescriptor<T> getOutputTypeDescriptor() {
return typeDescriptor;
}
+
+ // Deliberately make separate instances equal. The per-instance cache
must use identity
+ // because equal instances can have different generic types.
+ @Override
+ public boolean equals(@Nullable Object other) {
+ return other instanceof DynamicTypeDoFn;
+ }
+
+ @Override
+ public int hashCode() {
+ return 0;
+ }
}
DoFn<String, String> stringFn = new
DynamicTypeDoFn<>(TypeDescriptors.strings());
@@ -1456,4 +1469,38 @@ public class DoFnInvokersTest {
stringInvoker.getClass(),
intInvoker.getClass());
}
+
+ @Test
+ public void testTypeDescriptorResolutionIsCachedPerDoFnInstance() {
+ class CountingTypeDoFn extends DoFn<String, String> {
+ private int inputTypeDescriptorCalls;
+ private int outputTypeDescriptorCalls;
+
+ @ProcessElement
+ public void processElement(@Element String element,
OutputReceiver<String> out) {
+ out.output(element);
+ }
+
+ @Override
+ public TypeDescriptor<String> getInputTypeDescriptor() {
+ inputTypeDescriptorCalls++;
+ return TypeDescriptors.strings();
+ }
+
+ @Override
+ public TypeDescriptor<String> getOutputTypeDescriptor() {
+ outputTypeDescriptorCalls++;
+ return TypeDescriptors.strings();
+ }
+ }
+
+ CountingTypeDoFn fn = new CountingTypeDoFn();
+
+ DoFnInvoker<String, String> firstInvoker = DoFnInvokers.invokerFor(fn);
+ DoFnInvoker<String, String> secondInvoker = DoFnInvokers.invokerFor(fn);
+
+ assertSame(firstInvoker.getClass(), secondInvoker.getClass());
+ assertEquals(1, fn.inputTypeDescriptorCalls);
+ assertEquals(1, fn.outputTypeDescriptorCalls);
+ }
}
