This is an automated email from the ASF dual-hosted git repository.
kennknowles 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 d688e1b9505 [Java SDK] Warn when ValueState contains collection types
(#37530)
d688e1b9505 is described below
commit d688e1b950587f8043990742b2718a652718a55c
Author: ZIHAN DAI <[email protected]>
AuthorDate: Fri Jul 24 00:20:16 2026 +1000
[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.
---
.../sdk/transforms/reflect/DoFnSignatures.java | 49 +++++++++
.../sdk/transforms/reflect/DoFnSignaturesTest.java | 117 +++++++++++++++++++++
2 files changed, 166 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 6d28f934386..431ecea54b6 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
@@ -106,6 +106,8 @@ import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Immuta
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 @@ import org.joda.time.Instant;
})
public class DoFnSignatures {
+ private static final Logger LOG =
LoggerFactory.getLogger(DoFnSignatures.class);
+
private DoFnSignatures() {}
/**
@@ -2368,12 +2372,57 @@ public class DoFnSignatures {
(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 =
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 330bb5b9441..d7c42e42153 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.TimerSpec;
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 class DoFnSignaturesTest {
@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 class DoFnSignaturesTest {
@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");
+ }
}
