From d75ea9d65aba52d9a0b8fed62184efd61d45715d Mon Sep 17 00:00:00 2001 From: ffccites <99155080+PDGGK@users.noreply.github.com> Date: Fri, 6 Feb 2026 10:29:02 +0800 Subject: [PATCH 1/5] [Java SDK] Warn when ValueState contains collection types When users declare ValueState, ValueState, or ValueState, log a warning suggesting they use MapState, BagState, or SetState instead. Storing collections in ValueState requires reading and writing the entire collection on each access, which can cause performance issues for large collections. The specialized state types provide better performance. Fixes #36746 --- .../transforms/reflect/DoFnSignatures.java | 69 +++++++++++++++++++ .../reflect/DoFnSignaturesTest.java | 61 ++++++++++++++++ 2 files changed, 130 insertions(+) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java index 310736c014cc..bda1beb2b71c 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java @@ -104,6 +104,8 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** Utilities for working with {@link DoFnSignature}. See {@link #getSignature}. */ @Internal @@ -113,6 +115,8 @@ }) public class DoFnSignatures { + private static final Logger LOG = LoggerFactory.getLogger(DoFnSignatures.class); + private DoFnSignatures() {} /** @@ -2327,12 +2331,77 @@ private static Map analyzeStateDeclarati (TypeDescriptor) TypeDescriptor.of(fnClazz).resolveType(unresolvedStateType); + // Warn if ValueState contains a collection type that could benefit from specialized state + warnIfValueStateContainsCollection(fnClazz, id, stateType); + declarations.put(id, DoFnSignature.StateDeclaration.create(id, field, stateType)); } return ImmutableMap.copyOf(declarations); } + /** + * Warns if a ValueState is declared with a collection type (Map, List, Set) that could benefit + * from using specialized state types (MapState, BagState, SetState) for better performance. + */ + private static void warnIfValueStateContainsCollection( + Class fnClazz, String stateId, TypeDescriptor stateType) { + if (!stateType.isSubtypeOf(TypeDescriptor.of(ValueState.class))) { + return; + } + + try { + // Get the type directly and extract ValueState's type parameter + Type type = stateType.getType(); + if (!(type instanceof ParameterizedType)) { + return; + } + + // Find ValueState in the type hierarchy and get its type argument + Type valueType = null; + ParameterizedType pType = (ParameterizedType) type; + if (pType.getRawType() == ValueState.class) { + valueType = pType.getActualTypeArguments()[0]; + } else { + // For subtypes of ValueState, we need to resolve the type parameter + return; + } + + if (valueType == null + || valueType instanceof java.lang.reflect.TypeVariable + || valueType instanceof java.lang.reflect.WildcardType) { + // Cannot determine actual type, skip warning + return; + } + + TypeDescriptor valueTypeDescriptor = TypeDescriptor.of(valueType); + Class rawType = valueTypeDescriptor.getRawType(); + + String recommendation = null; + if (Map.class.isAssignableFrom(rawType)) { + recommendation = "MapState"; + } else if (List.class.isAssignableFrom(rawType)) { + recommendation = "BagState or OrderedListState"; + } else if (java.util.Set.class.isAssignableFrom(rawType)) { + recommendation = "SetState"; + } + + if (recommendation != null) { + LOG.warn( + "DoFn {} declares ValueState '{}' with type {}. " + + "Storing collections in ValueState requires reading and writing the entire " + + "collection on each access, which can cause performance issues. " + + "Consider using {} instead for better performance with large collections.", + fnClazz.getSimpleName(), + stateId, + rawType.getSimpleName(), + recommendation); + } + } catch (Exception e) { + // If we can't determine the type, don't warn - it's just an optimization hint + } + } + private static @Nullable Method findAnnotatedMethod( ErrorReporter errors, Class anno, Class fnClazz, boolean required) { Collection matches = diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java index de4a622e03d7..a394b23cd7a0 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java @@ -1700,4 +1700,65 @@ public void onMyTimer() {} @Override public void processWithTimer(ProcessContext context, Timer timer) {} } + + // Test DoFns for ValueState collection warning tests + private static class DoFnWithMapValueState extends DoFn { + @StateId("mapState") + private final StateSpec>> mapState = + StateSpecs.value(); + + @ProcessElement + public void process() {} + } + + private static class DoFnWithListValueState extends DoFn { + @StateId("listState") + private final StateSpec>> listState = StateSpecs.value(); + + @ProcessElement + public void process() {} + } + + private static class DoFnWithSetValueState extends DoFn { + @StateId("setState") + private final StateSpec>> setState = StateSpecs.value(); + + @ProcessElement + public void process() {} + } + + private static class DoFnWithSimpleValueState extends DoFn { + @StateId("simpleState") + private final StateSpec> simpleState = StateSpecs.value(); + + @ProcessElement + public void process() {} + } + + @Test + public void testValueStateWithMapLogsWarning() { + // This test verifies that the signature can be parsed for DoFns with collection ValueState. + // The warning is logged but doesn't prevent the signature from being created. + DoFnSignature signature = DoFnSignatures.getSignature(DoFnWithMapValueState.class); + assertThat(signature.stateDeclarations().get("mapState"), notNullValue()); + } + + @Test + public void testValueStateWithListLogsWarning() { + DoFnSignature signature = DoFnSignatures.getSignature(DoFnWithListValueState.class); + assertThat(signature.stateDeclarations().get("listState"), notNullValue()); + } + + @Test + public void testValueStateWithSetLogsWarning() { + DoFnSignature signature = DoFnSignatures.getSignature(DoFnWithSetValueState.class); + assertThat(signature.stateDeclarations().get("setState"), notNullValue()); + } + + @Test + public void testValueStateWithSimpleTypeNoWarning() { + // Simple types should not trigger any warning + DoFnSignature signature = DoFnSignatures.getSignature(DoFnWithSimpleValueState.class); + assertThat(signature.stateDeclarations().get("simpleState"), notNullValue()); + } } From 390f421ec6fb5d14079189f424671a05ed807b86 Mon Sep 17 00:00:00 2001 From: ffccites <99155080+PDGGK@users.noreply.github.com> Date: Sat, 7 Feb 2026 13:55:31 +0800 Subject: [PATCH 2/5] Refine ValueState collection hint: downgrade to INFO and clarify trade-offs - Change LOG.warn to LOG.info (performance hint, not correctness issue) - Clarify that ValueState is appropriate for small collections or atomic replacement - Add runner compatibility caveat for specialized state types - Address community feedback from @Eliaaazzz --- .../beam/sdk/transforms/reflect/DoFnSignatures.java | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java index bda1beb2b71c..96b83dc656e4 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java @@ -2387,11 +2387,12 @@ private static void warnIfValueStateContainsCollection( } if (recommendation != null) { - LOG.warn( - "DoFn {} declares ValueState '{}' with type {}. " - + "Storing collections in ValueState requires reading and writing the entire " - + "collection on each access, which can cause performance issues. " - + "Consider using {} instead for better performance with large collections.", + LOG.info( + "DoFn {} declares ValueState '{}' with collection type {}. " + + "ValueState reads/writes the entire collection on each access. " + + "This is appropriate for small collections or atomic replacement. " + + "For large collections or frequent appends, consider using {} instead " + + "(if supported by your runner).", fnClazz.getSimpleName(), stateId, rawType.getSimpleName(), From 877547d67d251cc973c496ab95f327af6e36df31 Mon Sep 17 00:00:00 2001 From: ffccites <99155080+PDGGK@users.noreply.github.com> Date: Sat, 7 Feb 2026 14:51:57 +0800 Subject: [PATCH 3/5] Address reviewer feedback: use TypeDescriptor API throughout - Use TypeDescriptor.resolveType() instead of raw Type manipulation - Use hasUnresolvedParameters() instead of instanceof checks - Use isSubtypeOf() for collection type detection - Remove catch-all Exception block entirely Addresses @kennknowles code review comments --- .../transforms/reflect/DoFnSignatures.java | 77 +++++++------------ 1 file changed, 29 insertions(+), 48 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java index 96b83dc656e4..03de49445e2d 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java @@ -2350,56 +2350,37 @@ private static void warnIfValueStateContainsCollection( return; } - try { - // Get the type directly and extract ValueState's type parameter - Type type = stateType.getType(); - if (!(type instanceof ParameterizedType)) { - return; - } - - // Find ValueState in the type hierarchy and get its type argument - Type valueType = null; - ParameterizedType pType = (ParameterizedType) type; - if (pType.getRawType() == ValueState.class) { - valueType = pType.getActualTypeArguments()[0]; - } else { - // For subtypes of ValueState, we need to resolve the type parameter - return; - } + // Use TypeDescriptor.resolveType() to extract ValueState's type parameter + // This preserves generic type information better than raw Type manipulation + TypeDescriptor valueTypeDescriptor = + stateType.resolveType(ValueState.class.getTypeParameters()[0]); - if (valueType == null - || valueType instanceof java.lang.reflect.TypeVariable - || valueType instanceof java.lang.reflect.WildcardType) { - // Cannot determine actual type, skip warning - return; - } - - TypeDescriptor valueTypeDescriptor = TypeDescriptor.of(valueType); - Class rawType = valueTypeDescriptor.getRawType(); - - String recommendation = null; - if (Map.class.isAssignableFrom(rawType)) { - recommendation = "MapState"; - } else if (List.class.isAssignableFrom(rawType)) { - recommendation = "BagState or OrderedListState"; - } else if (java.util.Set.class.isAssignableFrom(rawType)) { - recommendation = "SetState"; - } + // Skip if the type has unresolved parameters (e.g., TypeVariable, WildcardType) + if (valueTypeDescriptor.hasUnresolvedParameters()) { + return; + } - if (recommendation != null) { - LOG.info( - "DoFn {} declares ValueState '{}' with collection type {}. " - + "ValueState reads/writes the entire collection on each access. " - + "This is appropriate for small collections or atomic replacement. " - + "For large collections or frequent appends, consider using {} instead " - + "(if supported by your runner).", - fnClazz.getSimpleName(), - stateId, - rawType.getSimpleName(), - recommendation); - } - } catch (Exception e) { - // If we can't determine the type, don't warn - it's just an optimization hint + // Use TypeDescriptor.isSubtypeOf() for type checking - stays in TypeDescriptor API + String recommendation = null; + if (valueTypeDescriptor.isSubtypeOf(TypeDescriptor.of(Map.class))) { + recommendation = "MapState"; + } else if (valueTypeDescriptor.isSubtypeOf(TypeDescriptor.of(List.class))) { + recommendation = "BagState or OrderedListState"; + } else if (valueTypeDescriptor.isSubtypeOf(TypeDescriptor.of(java.util.Set.class))) { + recommendation = "SetState"; + } + + if (recommendation != null) { + LOG.info( + "DoFn {} declares ValueState '{}' with collection type {}. " + + "ValueState reads/writes the entire collection on each access. " + + "This is appropriate for small collections or atomic replacement. " + + "For large collections or frequent appends, consider using {} instead " + + "(if supported by your runner).", + fnClazz.getSimpleName(), + stateId, + valueTypeDescriptor.getRawType().getSimpleName(), + recommendation); } } From 04f6c5094c8c1f6cd9be1f913d281e7691f5a8c1 Mon Sep 17 00:00:00 2001 From: ffccites <99155080+PDGGK@users.noreply.github.com> Date: Sun, 22 Feb 2026 15:14:52 +0800 Subject: [PATCH 4/5] Upgrade collection hint from INFO to WARN per reviewer feedback --- .../apache/beam/sdk/transforms/reflect/DoFnSignatures.java | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java index 03de49445e2d..9421e1d1c9d2 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java @@ -2371,11 +2371,10 @@ private static void warnIfValueStateContainsCollection( } if (recommendation != null) { - LOG.info( + LOG.warn( "DoFn {} declares ValueState '{}' with collection type {}. " + "ValueState reads/writes the entire collection on each access. " - + "This is appropriate for small collections or atomic replacement. " - + "For large collections or frequent appends, consider using {} instead " + + "For large or frequently-updated collections, consider using {} instead " + "(if supported by your runner).", fnClazz.getSimpleName(), stateId, From ea57745ff91026d6c8beb68034dcc8eb2404b762 Mon Sep 17 00:00:00 2001 From: ffccites <99155080+PDGGK@users.noreply.github.com> Date: Fri, 26 Jun 2026 03:07:46 +1000 Subject: [PATCH 5/5] Warn for parameterized collection ValueState types ValueState>, ValueState> and ValueState> were silently skipped because the resolved value type has unresolved parameters (the element type is a type variable). The collection's own raw type is known in these cases, so the specialized-state recommendation (SetState / BagState or OrderedListState / MapState) can still be made. Drop the blanket hasUnresolvedParameters() early-return and match on the raw collection type instead; a payload that is itself a bare type variable (ValueState) still produces no recommendation. Add unit tests covering the parameterized Set/List cases and the bare type-variable no-warning case. --- .../transforms/reflect/DoFnSignatures.java | 11 ++-- .../reflect/DoFnSignaturesTest.java | 56 +++++++++++++++++++ 2 files changed, 61 insertions(+), 6 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java index 9421e1d1c9d2..bcccedc8d4d3 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java @@ -2355,12 +2355,11 @@ private static void warnIfValueStateContainsCollection( TypeDescriptor valueTypeDescriptor = stateType.resolveType(ValueState.class.getTypeParameters()[0]); - // Skip if the type has unresolved parameters (e.g., TypeVariable, WildcardType) - if (valueTypeDescriptor.hasUnresolvedParameters()) { - return; - } - - // Use TypeDescriptor.isSubtypeOf() for type checking - stays in TypeDescriptor API + // Match on the collection's raw type. We intentionally do not skip types with unresolved + // parameters: for a parameterized collection such as ValueState> the collection's own + // raw type (Set) is known even though the element type T is a type variable, so we can still + // recommend SetState. A payload that is itself a bare type variable (e.g. ValueState) simply + // matches none of the branches below and produces no recommendation. String recommendation = null; if (valueTypeDescriptor.isSubtypeOf(TypeDescriptor.of(Map.class))) { recommendation = "MapState"; diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java index a394b23cd7a0..3916d8689a2f 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java @@ -51,6 +51,7 @@ import org.apache.beam.sdk.state.TimerSpecs; import org.apache.beam.sdk.state.ValueState; import org.apache.beam.sdk.state.WatermarkHoldState; +import org.apache.beam.sdk.testing.ExpectedLogs; import org.apache.beam.sdk.testing.SerializableMatchers; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.Sum; @@ -101,6 +102,8 @@ public class DoFnSignaturesTest { @Rule public ExpectedException thrown = ExpectedException.none(); + @Rule public ExpectedLogs expectedLogs = ExpectedLogs.none(DoFnSignatures.class); + @Test public void testBasicDoFnProcessContext() throws Exception { DoFnSignature sig = @@ -1735,6 +1738,32 @@ private static class DoFnWithSimpleValueState extends DoFn { public void process() {} } + private static class DoFnWithParameterizedSetValueState extends DoFn { + @StateId("parameterizedSetState") + private final StateSpec>> parameterizedSetState = + StateSpecs.value(); + + @ProcessElement + public void process() {} + } + + private static class DoFnWithParameterizedListValueState extends DoFn { + @StateId("parameterizedListState") + private final StateSpec>> parameterizedListState = + StateSpecs.value(); + + @ProcessElement + public void process() {} + } + + private static class DoFnWithTypeVariableValueState extends DoFn { + @StateId("typeVariableState") + private final StateSpec> typeVariableState = StateSpecs.value(); + + @ProcessElement + public void process() {} + } + @Test public void testValueStateWithMapLogsWarning() { // This test verifies that the signature can be parsed for DoFns with collection ValueState. @@ -1760,5 +1789,32 @@ public void testValueStateWithSimpleTypeNoWarning() { // Simple types should not trigger any warning DoFnSignature signature = DoFnSignatures.getSignature(DoFnWithSimpleValueState.class); assertThat(signature.stateDeclarations().get("simpleState"), notNullValue()); + expectedLogs.verifyNotLogged("with collection type"); + } + + @Test + public void testValueStateWithParameterizedSetLogsWarning() { + // A parameterized collection such as ValueState> still has a known raw collection type + // (Set) even though the element type is a type variable, so it should still recommend SetState. + DoFnSignature signature = DoFnSignatures.getSignature(DoFnWithParameterizedSetValueState.class); + assertThat(signature.stateDeclarations().get("parameterizedSetState"), notNullValue()); + expectedLogs.verifyWarn("SetState"); + } + + @Test + public void testValueStateWithParameterizedListLogsWarning() { + DoFnSignature signature = + DoFnSignatures.getSignature(DoFnWithParameterizedListValueState.class); + assertThat(signature.stateDeclarations().get("parameterizedListState"), notNullValue()); + expectedLogs.verifyWarn("BagState or OrderedListState"); + } + + @Test + public void testValueStateWithBareTypeVariableNoWarning() { + // A bare type variable payload (ValueState) has no raw collection type to inspect and must + // not produce a collection recommendation. + DoFnSignature signature = DoFnSignatures.getSignature(DoFnWithTypeVariableValueState.class); + assertThat(signature.stateDeclarations().get("typeVariableState"), notNullValue()); + expectedLogs.verifyNotLogged("with collection type"); } }