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..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 @@ -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,57 @@ 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; + } + + // 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> 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"; + } 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 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..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 = @@ -1700,4 +1703,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 { + @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() {} + } + + 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. + // 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> 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"); + } }