[Java SDK] Warn when ValueState contains collection types (#37530)

When users declare ValueState<Map>, ValueState<List>, or ValueState<Set>,
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.

Author: PDGGK

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java

@@ -106,6 +106,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
@@ -115,6 +117,8 @@
 })
 public class DoFnSignatures {
 
+  private static final Logger LOG = LoggerFactory.getLogger(DoFnSignatures.class);
+
   private DoFnSignatures() {}
 
   /**
@@ -2368,12 +2372,57 @@ private static Map<String, DoFnSignature.StateDeclaration> analyzeStateDeclarati
           (TypeDescriptor<? extends State>)
               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<? extends State> stateType) {
+    if (!stateType.isSubtypeOf(TypeDescriptor.of(ValueState.class))) {
+      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]);
+
+    // Match on the collection's raw type. We intentionally do not skip types with unresolved
+    // parameters: for a parameterized collection such as ValueState<Set<T>> 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<T>) simply
+    // matches none of the branches below and produces no recommendation.
+    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.warn(
+          "DoFn {} declares ValueState '{}' with collection type {}. "
+              + "ValueState reads/writes the entire collection on each access. "
+              + "For large or frequently-updated collections, consider using {} instead "
+              + "(if supported by your runner).",
+          fnClazz.getSimpleName(),
+          stateId,
+          valueTypeDescriptor.getRawType().getSimpleName(),
+          recommendation);
+    }
+  }
+
   private static @Nullable Method findAnnotatedMethod(
       ErrorReporter errors, Class<? extends Annotation> anno, Class<?> fnClazz, boolean required) {
     Collection<Method> matches =

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;
@@ -137,6 +138,8 @@ public void process(@FireTimestamp Instant ts) {}
 
   @Rule public ExpectedException thrown = ExpectedException.none();
 
+  @Rule public ExpectedLogs expectedLogs = ExpectedLogs.none(DoFnSignatures.class);
+
   @Test
   public void testBasicDoFnProcessContext() throws Exception {
     DoFnSignature sig =
@@ -1810,4 +1813,118 @@ public void onMyTimer() {}
     @Override
     public void processWithTimer(ProcessContext context, Timer timer) {}
   }
+
+  // Test DoFns for ValueState collection warning tests
+  private static class DoFnWithMapValueState extends DoFn<String, String> {
+    @StateId("mapState")
+    private final StateSpec<ValueState<java.util.Map<String, String>>> mapState =
+        StateSpecs.value();
+
+    @ProcessElement
+    public void process() {}
+  }
+
+  private static class DoFnWithListValueState extends DoFn<String, String> {
+    @StateId("listState")
+    private final StateSpec<ValueState<java.util.List<String>>> listState = StateSpecs.value();
+
+    @ProcessElement
+    public void process() {}
+  }
+
+  private static class DoFnWithSetValueState extends DoFn<String, String> {
+    @StateId("setState")
+    private final StateSpec<ValueState<java.util.Set<String>>> setState = StateSpecs.value();
+
+    @ProcessElement
+    public void process() {}
+  }
+
+  private static class DoFnWithSimpleValueState extends DoFn<String, String> {
+    @StateId("simpleState")
+    private final StateSpec<ValueState<String>> simpleState = StateSpecs.value();
+
+    @ProcessElement
+    public void process() {}
+  }
+
+  private static class DoFnWithParameterizedSetValueState<T> extends DoFn<String, String> {
+    @StateId("parameterizedSetState")
+    private final StateSpec<ValueState<java.util.Set<T>>> parameterizedSetState =
+        StateSpecs.value();
+
+    @ProcessElement
+    public void process() {}
+  }
+
+  private static class DoFnWithParameterizedListValueState<T> extends DoFn<String, String> {
+    @StateId("parameterizedListState")
+    private final StateSpec<ValueState<java.util.List<T>>> parameterizedListState =
+        StateSpecs.value();
+
+    @ProcessElement
+    public void process() {}
+  }
+
+  private static class DoFnWithTypeVariableValueState<T> extends DoFn<String, String> {
+    @StateId("typeVariableState")
+    private final StateSpec<ValueState<T>> 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.
+    // 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());
+    expectedLogs.verifyNotLogged("with collection type");
+  }
+
+  @Test
+  public void testValueStateWithParameterizedSetLogsWarning() {
+    // A parameterized collection such as ValueState<Set<T>> 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<T>) 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");
+  }
 }