Skip to content

Fix OnWindowExpirationContext. #38878

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
reuvenlax merged 3 commits into from
Jun 9, 2026
Merged

Fix OnWindowExpirationContext. #38878

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
23f4b5d
foo
reuvenlax Jun 9, 2026
842a1fd
fix compilation
reuvenlax Jun 9, 2026
e8e7217
fix compilation
reuvenlax Jun 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -649,6 +649,13 @@ public DoFn<InputT, OutputT>.OnTimerContext onTimerContext(DoFn<InputT, OutputT>
"Cannot access OnTimerContext outside of @OnTimer methods.");
}

@Override
public DoFn<InputT, OutputT>.OnWindowExpirationContext onWindowExpirationContext(
DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
"Cannot access OnWindowExpirationContext outside of @OnWindowExpiration methods.");
}

@Override
public RestrictionTracker<?, ?> restrictionTracker() {
throw new UnsupportedOperationException("RestrictionTracker parameters are not supported.");
Expand Down Expand Up @@ -958,6 +965,13 @@ public DoFn<InputT, OutputT>.OnTimerContext onTimerContext(DoFn<InputT, OutputT>
return this;
}

@Override
public DoFn<InputT, OutputT>.OnWindowExpirationContext onWindowExpirationContext(
DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
"Cannot access OnWindowExpirationContext outside of @OnWindowExpiration methods.");
}

@Override
public RestrictionTracker<?, ?> restrictionTracker() {
throw new UnsupportedOperationException("RestrictionTracker parameters are not supported.");
Expand Down Expand Up @@ -1299,6 +1313,12 @@ public DoFn<InputT, OutputT>.OnTimerContext onTimerContext(DoFn<InputT, OutputT>
throw new UnsupportedOperationException("OnTimerContext parameters are not supported.");
}

@Override
public DoFn<InputT, OutputT>.OnWindowExpirationContext onWindowExpirationContext(
DoFn<InputT, OutputT> doFn) {
return this;
}

@Override
public RestrictionTracker<?, ?> restrictionTracker() {
throw new UnsupportedOperationException("RestrictionTracker parameters are not supported.");
Expand Down
12 changes: 12 additions & 0 deletions ...ore/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,8 @@ class ByteBuddyDoFnInvokerFactory implements DoFnInvokerFactory {
public static final String OUTPUT_PARAMETER_METHOD = "outputReceiver";
public static final String TAGGED_OUTPUT_PARAMETER_METHOD = "taggedOutputReceiver";
public static final String ON_TIMER_CONTEXT_PARAMETER_METHOD = "onTimerContext";
public static final String ON_WINDOW_EXPIRATION_CONTEXT_PARAMETER_METHOD =
"onWindowExpirationContext";
public static final String WINDOW_PARAMETER_METHOD = "window";
public static final String PANE_INFO_PARAMETER_METHOD = "paneInfo";
public static final String PIPELINE_OPTIONS_PARAMETER_METHOD = "pipelineOptions";
Expand Down Expand Up @@ -1170,6 +1172,16 @@ public StackManipulation dispatch(OnTimerContextParameter p) {
ON_TIMER_CONTEXT_PARAMETER_METHOD, DoFn.class)));
}

@Override
public StackManipulation dispatch(
DoFnSignature.Parameter.OnWindowExpirationContextParameter p) {
return new StackManipulation.Compound(
pushDelegate,
MethodInvocation.invoke(
getExtraContextFactoryMethodDescription(
ON_WINDOW_EXPIRATION_CONTEXT_PARAMETER_METHOD, DoFn.class)));
}

@Override
public StackManipulation dispatch(WindowParameter p) {
return new StackManipulation.Compound(
Expand Down
17 changes: 17 additions & 0 deletions sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,10 @@ interface ArgumentProvider<InputT, OutputT> {
/** Provide a {@link DoFn.OnTimerContext} to use with the given {@link DoFn}. */
DoFn<InputT, OutputT>.OnTimerContext onTimerContext(DoFn<InputT, OutputT> doFn);

/** Provide a {@link DoFn.OnWindowExpirationContext} to use with the given {@link DoFn}. */
DoFn<InputT, OutputT>.OnWindowExpirationContext onWindowExpirationContext(
DoFn<InputT, OutputT> doFn);

/** Provide a reference to the input element. */
InputT element(DoFn<InputT, OutputT> doFn);

Expand Down Expand Up @@ -447,6 +451,13 @@ public DoFn<InputT, OutputT>.OnTimerContext onTimerContext(DoFn<InputT, OutputT>
String.format("OnTimerContext unsupported in %s", getErrorContext()));
}

@Override
public DoFn<InputT, OutputT>.OnWindowExpirationContext onWindowExpirationContext(
DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
String.format("OnWindowExpirationContext unsupported in %s", getErrorContext()));
}

@Override
public State state(String stateId, boolean alwaysFetched) {
throw new UnsupportedOperationException(
Expand Down Expand Up @@ -538,6 +549,12 @@ public DoFn<InputT, OutputT>.OnTimerContext onTimerContext(DoFn<InputT, OutputT>
return delegate.onTimerContext(doFn);
}

@Override
public DoFn<InputT, OutputT>.OnWindowExpirationContext onWindowExpirationContext(
DoFn<InputT, OutputT> doFn) {
return delegate.onWindowExpirationContext(doFn);
}

@Override
public InputT element(DoFn<InputT, OutputT> doFn) {
return delegate.element(doFn);
Expand Down
9 changes: 9 additions & 0 deletions sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignature.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -305,6 +305,8 @@ public <ResultT> ResultT match(Cases<ResultT> cases) {
return cases.dispatch((ProcessContextParameter) this);
} else if (this instanceof OnTimerContextParameter) {
return cases.dispatch((OnTimerContextParameter) this);
} else if (this instanceof OnWindowExpirationContextParameter) {
return cases.dispatch((OnWindowExpirationContextParameter) this);
} else if (this instanceof WindowParameter) {
return cases.dispatch((WindowParameter) this);
} else if (this instanceof PaneInfoParameter) {
Expand Down Expand Up @@ -391,6 +393,8 @@ public interface Cases<ResultT> {

ResultT dispatch(OnTimerContextParameter p);

ResultT dispatch(OnWindowExpirationContextParameter p);

ResultT dispatch(WindowParameter p);

ResultT dispatch(PaneInfoParameter p);
Expand Down Expand Up @@ -498,6 +502,11 @@ public ResultT dispatch(OnTimerContextParameter p) {
return dispatchDefault(p);
}

@Override
public ResultT dispatch(OnWindowExpirationContextParameter p) {
return dispatchDefault(p);
}

@Override
public ResultT dispatch(WindowParameter p) {
return dispatchDefault(p);
Expand Down
3 changes: 2 additions & 1 deletion sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,7 +229,8 @@ private DoFnSignatures() {}
Parameter.StateParameter.class,
Parameter.TimestampParameter.class,
Parameter.KeyParameter.class,
Parameter.SideInputParameter.class);
Parameter.SideInputParameter.class,
Parameter.OnWindowExpirationContextParameter.class);

private static final Collection<Class<? extends Parameter>>
ALLOWED_GET_INITIAL_RESTRICTION_PARAMETERS =
Expand Down
6 changes: 6 additions & 0 deletions ...core/src/main/java/org/apache/beam/sdk/util/construction/SplittableParDoNaiveBounded.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -506,6 +506,12 @@ public DoFn<InputT, OutputT>.OnTimerContext onTimerContext(DoFn<InputT, OutputT>
throw new IllegalStateException();
}

@Override
public DoFn<InputT, OutputT>.OnWindowExpirationContext onWindowExpirationContext(
DoFn<InputT, OutputT> doFn) {
throw new IllegalStateException();
}

@Override
public InputT element(DoFn<InputT, OutputT> doFn) {
return element;
Expand Down
4 changes: 4 additions & 0 deletions sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7335,11 +7335,15 @@ public void onTimer(
@OnWindowExpiration
public void onWindowExpiration(
@AlwaysFetched @StateId(stateId) ValueState<Integer> state,
BoundedWindow window,
@Key String key,
OnWindowExpirationContext context,
OutputReceiver<Integer> r) {
Integer currentValue = MoreObjects.firstNonNull(state.read(), 0);
// verify state
assertEquals(1, (int) currentValue);
Preconditions.checkNotNull(context);

@gemini-code-assist gemini-code-assist Bot Jun 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

In unit tests, it is more idiomatic to use JUnit's assertNotNull instead of Guava's Preconditions.checkNotNull. Preconditions.checkNotNull throws a NullPointerException (which JUnit treats as an unexpected error), whereas assertNotNull throws an AssertionError (which JUnit treats as a test failure), providing better integration with test runners and clearer test results.

Suggested change
Preconditions.checkNotNull(context);
assertNotNull(context);

assertEquals(window, context.window());
// To check output is received from OnWindowExpiration
r.output(currentValue);
}
Expand Down
8 changes: 6 additions & 2 deletions sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1422,16 +1422,20 @@ public void bar(
@StateId("foo") ValueState<Integer> s,
PipelineOptions p,
OutputReceiver<String> o,
MultiOutputReceiver m) {}
MultiOutputReceiver m,
OnWindowExpirationContext c) {}
}.getClass());

List<Parameter> params = sig.onWindowExpiration().extraParameters();
assertThat(params.size(), equalTo(5));
assertThat(params.size(), equalTo(6));
assertThat(params.get(0), instanceOf(WindowParameter.class));
assertThat(params.get(1), instanceOf(StateParameter.class));
assertThat(params.get(2), instanceOf(PipelineOptionsParameter.class));
assertThat(params.get(3), instanceOf(OutputReceiverParameter.class));
assertThat(params.get(4), instanceOf(TaggedOutputReceiverParameter.class));
assertThat(
params.get(5),
instanceOf(DoFnSignature.Parameter.OnWindowExpirationContextParameter.class));
}

private interface FeatureTest {
Expand Down
13 changes: 13 additions & 0 deletions sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2250,6 +2250,13 @@ public DoFn<InputT, OutputT>.OnTimerContext onTimerContext(DoFn<InputT, OutputT>
"Cannot access OnTimerContext outside of @OnTimer methods.");
}

@Override
public DoFn<InputT, OutputT>.OnWindowExpirationContext onWindowExpirationContext(
DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
"Cannot access OnWindowExpirationContext outside of @OnWindowExpiration methods.");
}

@Override
public RestrictionTracker<?, ?> restrictionTracker() {
return currentTracker;
Expand Down Expand Up @@ -2469,6 +2476,12 @@ private void checkOnWindowExpirationTimestamp(Instant timestamp) {
private final OnWindowExpirationContext.Context context =
new OnWindowExpirationContext.Context();

@Override
public DoFn<InputT, OutputT>.OnWindowExpirationContext onWindowExpirationContext(
DoFn<InputT, OutputT> doFn) {
return context;
}

@Override
public BoundedWindow window() {
return currentWindow;
Expand Down
Loading