diff --git a/CHANGES.md b/CHANGES.md index 81f7d493677b..4eb1db840d6f 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -103,6 +103,7 @@ ## New Features / Improvements +* Added automatic caching of bounded, single-pane side-input views for classic Java Flink DataStream execution ([#39866](https://github.com/apache/beam/issues/39866)). * Added `GroupIntoBatches` transform and the standard `beam:coder:sharded_key:v1` coder to the Go SDK, along with `beam.Coder.IsDeterministic`, `beam.PCollection.WindowingStrategy`, diff --git a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/CachedSideInputReader.java b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/CachedSideInputReader.java new file mode 100644 index 000000000000..ee0d9df53b57 --- /dev/null +++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/CachedSideInputReader.java @@ -0,0 +1,93 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.flink.translation.wrappers.streaming; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.DefaultTrigger; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.flink.api.common.JobID; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; + +/** {@link SideInputReader} that caches single-pane materialized views within a TaskManager JVM. */ +public final class CachedSideInputReader implements SideInputReader { + + public static CachedSideInputReader of( + JobID jobId, + int attemptNumber, + SideInputReader delegate, + Collection> cacheableViews) { + return new CachedSideInputReader(jobId, attemptNumber, delegate, cacheableViews); + } + + static Collection> cacheableViews(Collection> sideInputs) { + Collection> cacheableViews = new ArrayList<>(); + for (PCollectionView view : sideInputs) { + PCollection pCollection = view.getPCollection(); + WindowingStrategy strategy = view.getWindowingStrategyInternal(); + if (pCollection != null + && pCollection.isBounded() == PCollection.IsBounded.BOUNDED + && strategy.getTrigger() instanceof DefaultTrigger + && Duration.ZERO.equals(strategy.getAllowedLateness())) { + cacheableViews.add(view); + } + } + return Collections.unmodifiableCollection(cacheableViews); + } + + private final JobID jobId; + private final int attemptNumber; + private final SideInputReader delegate; + private final Collection> cacheableViews; + + private CachedSideInputReader( + JobID jobId, + int attemptNumber, + SideInputReader delegate, + Collection> cacheableViews) { + this.jobId = jobId; + this.attemptNumber = attemptNumber; + this.delegate = delegate; + this.cacheableViews = cacheableViews; + } + + @Override + public @Nullable T get(PCollectionView view, BoundedWindow window) { + if (!cacheableViews.contains(view)) { + return delegate.get(view, window); + } + return SideInputCache.getOrMaterialize( + jobId, attemptNumber, view, window, () -> delegate.get(view, window)); + } + + @Override + public boolean contains(PCollectionView view) { + return delegate.contains(view); + } + + @Override + public boolean isEmpty() { + return delegate.isEmpty(); + } +} diff --git a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java index e582635e0988..c0d3fc21307c 100644 --- a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java +++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java @@ -96,6 +96,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.flink.api.common.JobID; import org.apache.flink.api.common.operators.ProcessingTimeService.ProcessingTimeCallback; import org.apache.flink.api.common.state.ListState; import org.apache.flink.api.common.state.ListStateDescriptor; @@ -162,6 +163,7 @@ public class DoFnOperator protected final List> additionalOutputTags; protected final Collection> sideInputs; + private final Collection> cacheableSideInputs; protected final Map> sideInputTagMapping; protected final WindowingStrategy windowingStrategy; @@ -297,6 +299,7 @@ public DoFnOperator( this.additionalOutputTags = additionalOutputTags; this.sideInputTagMapping = sideInputTagMapping; this.sideInputs = sideInputs; + this.cacheableSideInputs = CachedSideInputReader.cacheableViews(sideInputs); this.serializedOptions = new SerializablePipelineOptions(options); this.isStreaming = serializedOptions.get().as(FlinkPipelineOptions.class).isStreaming(); this.windowingStrategy = windowingStrategy; @@ -466,7 +469,12 @@ public void initializeState(StateInitializationContext context) throws Exception serializedOptions); sideInputHandler = new SideInputHandler(sideInputs, sideInputStateInternals); - sideInputReader = sideInputHandler; + sideInputReader = + createSideInputReader( + cacheableSideInputs, + getContainingTask().getEnvironment().getJobID(), + getContainingTask().getEnvironment().getTaskInfo().getAttemptNumber(), + sideInputHandler); Stream> pushedBack = pushedBackElementsHandler.getElements(); long min = @@ -790,6 +798,27 @@ protected void addSideInputValue(StreamRecord streamRecord) { PCollectionView sideInput = sideInputTagMapping.get(streamRecord.getValue().getUnionTag()); sideInputHandler.addSideInputValue(sideInput, value); + // Invalidate only after the state write: a concurrent reader that re-caches between an + // earlier invalidation and the write would pin the previous value with no later invalidation. + for (BoundedWindow window : value.getWindows()) { + SideInputCache.invalidate( + getContainingTask().getEnvironment().getJobID(), + getContainingTask().getEnvironment().getTaskInfo().getAttemptNumber(), + sideInput, + window); + } + } + + @VisibleForTesting + static SideInputReader createSideInputReader( + Collection> cacheableViews, + JobID jobId, + int attemptNumber, + SideInputReader delegate) { + if (!cacheableViews.isEmpty()) { + return CachedSideInputReader.of(jobId, attemptNumber, delegate, cacheableViews); + } + return delegate; } @Override diff --git a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SideInputCache.java b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SideInputCache.java new file mode 100644 index 000000000000..ce49d8a5de99 --- /dev/null +++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SideInputCache.java @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.flink.translation.wrappers.streaming; + +import java.util.Objects; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.UncheckedExecutionException; +import org.apache.flink.api.common.JobID; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** Process-wide cache of materialized side-input views. */ +final class SideInputCache { + + // Materialized view sizes are unknown to the runner, so the cache cannot be bounded by weight; + // soft values let the JVM reclaim entries under memory pressure instead of failing with OOM. + // Operator instances must not clear job entries when they close because peer subtasks and later + // operators can still use them. Expiration bounds entries after their last access. + // Attempt-specific keys prevent restored state from using a value cached by an earlier attempt. + private static final Cache, Value> MATERIALIZED_SIDE_INPUTS = + CacheBuilder.newBuilder().expireAfterAccess(5, TimeUnit.MINUTES).softValues().build(); + + private SideInputCache() {} + + static @Nullable T getOrMaterialize( + JobID jobId, + int attemptNumber, + PCollectionView view, + BoundedWindow window, + Supplier<@Nullable T> materializer) { + @SuppressWarnings("unchecked") + Cache, Value> cache = + (Cache, Value>) (Cache) MATERIALIZED_SIDE_INPUTS; + try { + return cache + .get(new Key<>(jobId, attemptNumber, view, window), () -> new Value<>(materializer.get())) + .getValue(); + } catch (ExecutionException | UncheckedExecutionException e) { + Throwable cause = e.getCause() != null ? e.getCause() : e; + Throwables.throwIfUnchecked(cause); + throw new RuntimeException(cause); + } + } + + static void invalidate( + JobID jobId, int attemptNumber, PCollectionView view, BoundedWindow window) { + MATERIALIZED_SIDE_INPUTS.invalidate(new Key<>(jobId, attemptNumber, view, window)); + } + + private static final class Key { + private final JobID jobId; + private final int attemptNumber; + private final PCollectionView view; + private final BoundedWindow window; + + private Key(JobID jobId, int attemptNumber, PCollectionView view, BoundedWindow window) { + this.jobId = jobId; + this.attemptNumber = attemptNumber; + this.view = view; + this.window = window; + } + + @Override + public boolean equals(@Nullable Object object) { + if (this == object) { + return true; + } + if (!(object instanceof Key)) { + return false; + } + Key other = (Key) object; + return Objects.equals(jobId, other.jobId) + && attemptNumber == other.attemptNumber + && Objects.equals(view, other.view) + && Objects.equals(window, other.window); + } + + @Override + public int hashCode() { + return Objects.hash(jobId, attemptNumber, view, window); + } + } + + /** Guava caches reject null values, but null is valid for a side-input reader. */ + private static final class Value { + private final @Nullable T value; + + private Value(@Nullable T value) { + this.value = value; + } + + private @Nullable T getValue() { + return value; + } + } +} diff --git a/runners/flink/2.0/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java b/runners/flink/2.0/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java index f1e35fafe83b..f3cc919a7d44 100644 --- a/runners/flink/2.0/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java +++ b/runners/flink/2.0/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java @@ -97,7 +97,6 @@ public void testDefaults() { assertThat(options.getAllowNonRestoredState(), is(false)); assertThat(options.getDisableMetrics(), is(false)); assertThat(options.getFasterCopy(), is(false)); - assertThat(options.isStreaming(), is(false)); assertThat(options.getMaxBundleSize(), is(5000L)); assertThat(options.getMaxBundleTimeMills(), is(10000L)); diff --git a/runners/flink/2.0/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkCachedSideInputReaderTest.java b/runners/flink/2.0/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkCachedSideInputReaderTest.java new file mode 100644 index 000000000000..c71dcd295449 --- /dev/null +++ b/runners/flink/2.0/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkCachedSideInputReaderTest.java @@ -0,0 +1,283 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.flink.translation.wrappers.streaming; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.nullValue; +import static org.junit.Assert.assertThrows; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.DefaultTrigger; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.transforms.windowing.IntervalWindow; +import org.apache.beam.sdk.transforms.windowing.Trigger; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.flink.api.common.JobID; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.Test; + +/** Tests for cached materialization of Flink side-input views. */ +public class FlinkCachedSideInputReaderTest { + + private static final int INITIAL_ATTEMPT = 0; + private static final int RETRY_ATTEMPT = 1; + + @Test + public void repeatedGetMaterializesOnce() { + JobID jobId = new JobID(); + PCollectionView view = view(); + CountingSideInputReader delegate = new CountingSideInputReader("value"); + SideInputReader reader = + CachedSideInputReader.of(jobId, INITIAL_ATTEMPT, delegate, Collections.singleton(view)); + + assertThat(reader.get(view, GlobalWindow.INSTANCE), is("value")); + assertThat(reader.get(view, GlobalWindow.INSTANCE), is("value")); + assertThat(delegate.getCount(), is(1)); + } + + @Test + public void readerInstancesForSameJobShareMaterialization() { + JobID jobId = new JobID(); + PCollectionView view = view(); + CountingSideInputReader delegate = new CountingSideInputReader("value"); + Collection> views = Collections.singleton(view); + + CachedSideInputReader.of(jobId, INITIAL_ATTEMPT, delegate, views) + .get(view, GlobalWindow.INSTANCE); + CachedSideInputReader.of(jobId, INITIAL_ATTEMPT, delegate, views) + .get(view, GlobalWindow.INSTANCE); + + assertThat(delegate.getCount(), is(1)); + } + + @Test + public void retryAttemptRematerializesValue() { + JobID jobId = new JobID(); + PCollectionView view = view(); + CountingSideInputReader delegate = new CountingSideInputReader("value"); + Collection> views = Collections.singleton(view); + + CachedSideInputReader.of(jobId, INITIAL_ATTEMPT, delegate, views) + .get(view, GlobalWindow.INSTANCE); + CachedSideInputReader.of(jobId, RETRY_ATTEMPT, delegate, views) + .get(view, GlobalWindow.INSTANCE); + + assertThat(delegate.getCount(), is(2)); + } + + @Test + public void keyIncludesViewWindowAndJob() { + PCollectionView firstView = view(); + PCollectionView secondView = view(); + IntervalWindow firstWindow = new IntervalWindow(Instant.EPOCH, Instant.ofEpochMilli(10)); + IntervalWindow secondWindow = + new IntervalWindow(Instant.ofEpochMilli(10), Instant.ofEpochMilli(20)); + CountingSideInputReader delegate = new CountingSideInputReader("value"); + JobID firstJob = new JobID(); + + Collection> views = Arrays.asList(firstView, secondView); + CachedSideInputReader.of(firstJob, INITIAL_ATTEMPT, delegate, views) + .get(firstView, firstWindow); + CachedSideInputReader.of(firstJob, INITIAL_ATTEMPT, delegate, views) + .get(secondView, firstWindow); + CachedSideInputReader.of(firstJob, INITIAL_ATTEMPT, delegate, views) + .get(firstView, secondWindow); + CachedSideInputReader.of(new JobID(), INITIAL_ATTEMPT, delegate, views) + .get(firstView, firstWindow); + + assertThat(delegate.getCount(), is(4)); + } + + @Test + public void invalidateRematerializesValue() { + JobID jobId = new JobID(); + PCollectionView view = view(); + CountingSideInputReader delegate = new CountingSideInputReader("value"); + SideInputReader reader = + CachedSideInputReader.of(jobId, INITIAL_ATTEMPT, delegate, Collections.singleton(view)); + + reader.get(view, GlobalWindow.INSTANCE); + SideInputCache.invalidate(jobId, INITIAL_ATTEMPT, view, GlobalWindow.INSTANCE); + reader.get(view, GlobalWindow.INSTANCE); + + assertThat(delegate.getCount(), is(2)); + } + + @Test + public void cachesNull() { + JobID jobId = new JobID(); + PCollectionView view = view(); + CountingSideInputReader delegate = new CountingSideInputReader(null); + SideInputReader reader = + CachedSideInputReader.of(jobId, INITIAL_ATTEMPT, delegate, Collections.singleton(view)); + + assertThat(reader.get(view, GlobalWindow.INSTANCE), nullValue()); + assertThat(reader.get(view, GlobalWindow.INSTANCE), nullValue()); + assertThat(delegate.getCount(), is(1)); + } + + @Test + public void automaticallyWrapsReaderWithCacheableViews() { + JobID jobId = new JobID(); + SideInputReader delegate = new CountingSideInputReader("value"); + PCollectionView view = cacheableView(); + Collection> cacheableViews = + CachedSideInputReader.cacheableViews(Collections.singleton(view)); + + assertThat( + DoFnOperator.createSideInputReader( + Collections.emptyList(), jobId, INITIAL_ATTEMPT, delegate), + is(delegate)); + + assertThat( + DoFnOperator.createSideInputReader(cacheableViews, jobId, INITIAL_ATTEMPT, delegate), + instanceOf(CachedSideInputReader.class)); + } + + @Test + public void selectsOnlyBoundedDefaultTriggerViewsWithoutLateness() { + PCollectionView cacheableView = cacheableView(); + PCollectionView unboundedView = + view(PCollection.IsBounded.UNBOUNDED, DefaultTrigger.of(), Duration.ZERO); + PCollectionView customTriggerView = + view(PCollection.IsBounded.BOUNDED, mock(Trigger.class), Duration.ZERO); + PCollectionView lateDataView = + view(PCollection.IsBounded.BOUNDED, DefaultTrigger.of(), Duration.standardMinutes(1)); + + Collection> cacheableViews = + CachedSideInputReader.cacheableViews( + Arrays.asList(cacheableView, unboundedView, customTriggerView, lateDataView)); + + assertThat(cacheableViews.size(), is(1)); + assertThat(cacheableViews.contains(cacheableView), is(true)); + } + + @Test + public void nonCacheableViewAlwaysUsesDelegate() { + PCollectionView cacheableView = view(); + PCollectionView nonCacheableView = view(); + CountingSideInputReader delegate = new CountingSideInputReader("value"); + SideInputReader reader = + CachedSideInputReader.of( + new JobID(), INITIAL_ATTEMPT, delegate, Collections.singleton(cacheableView)); + + reader.get(cacheableView, GlobalWindow.INSTANCE); + reader.get(cacheableView, GlobalWindow.INSTANCE); + reader.get(nonCacheableView, GlobalWindow.INSTANCE); + reader.get(nonCacheableView, GlobalWindow.INSTANCE); + + assertThat(delegate.getCount(), is(3)); + } + + @Test + public void materializationExceptionPropagatesUnwrapped() { + PCollectionView view = view(); + SideInputReader reader = + CachedSideInputReader.of( + new JobID(), + INITIAL_ATTEMPT, + new SideInputReader() { + @Override + public @Nullable T get(PCollectionView view, BoundedWindow window) { + throw new IllegalStateException("materialization failed"); + } + + @Override + public boolean contains(PCollectionView view) { + return true; + } + + @Override + public boolean isEmpty() { + return false; + } + }, + Collections.singleton(view)); + + IllegalStateException exception = + assertThrows(IllegalStateException.class, () -> reader.get(view, GlobalWindow.INSTANCE)); + assertThat(exception.getMessage(), is("materialization failed")); + } + + private static PCollectionView cacheableView() { + return view(PCollection.IsBounded.BOUNDED, DefaultTrigger.of(), Duration.ZERO); + } + + @SuppressWarnings("unchecked") + private static PCollectionView view( + PCollection.IsBounded bounded, Trigger trigger, Duration allowedLateness) { + PCollectionView view = mock(PCollectionView.class); + PCollection pCollection = mock(PCollection.class); + WindowingStrategy strategy = mock(WindowingStrategy.class); + doReturn(pCollection).when(view).getPCollection(); + when(pCollection.isBounded()).thenReturn(bounded); + doReturn(strategy).when(view).getWindowingStrategyInternal(); + when(strategy.getTrigger()).thenReturn(trigger); + when(strategy.getAllowedLateness()).thenReturn(allowedLateness); + return view; + } + + @SuppressWarnings("unchecked") + private static PCollectionView view() { + return mock(PCollectionView.class); + } + + private static final class CountingSideInputReader implements SideInputReader { + private final AtomicInteger getCount = new AtomicInteger(); + private final @Nullable Object value; + + private CountingSideInputReader(@Nullable Object value) { + this.value = value; + } + + @Override + @SuppressWarnings("unchecked") + public @Nullable T get(PCollectionView view, BoundedWindow window) { + getCount.incrementAndGet(); + return (T) value; + } + + @Override + public boolean contains(PCollectionView view) { + return true; + } + + @Override + public boolean isEmpty() { + return false; + } + + private int getCount() { + return getCount.get(); + } + } +} diff --git a/runners/flink/2.2/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java b/runners/flink/2.2/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java index 409797625db4..752100c97722 100644 --- a/runners/flink/2.2/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java +++ b/runners/flink/2.2/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java @@ -96,6 +96,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.flink.api.common.JobID; import org.apache.flink.api.common.operators.ProcessingTimeService.ProcessingTimeCallback; import org.apache.flink.api.common.state.ListState; import org.apache.flink.api.common.state.ListStateDescriptor; @@ -162,6 +163,7 @@ public class DoFnOperator protected final List> additionalOutputTags; protected final Collection> sideInputs; + private final Collection> cacheableSideInputs; protected final Map> sideInputTagMapping; protected final WindowingStrategy windowingStrategy; @@ -297,6 +299,7 @@ public DoFnOperator( this.additionalOutputTags = additionalOutputTags; this.sideInputTagMapping = sideInputTagMapping; this.sideInputs = sideInputs; + this.cacheableSideInputs = CachedSideInputReader.cacheableViews(sideInputs); this.serializedOptions = new SerializablePipelineOptions(options); this.isStreaming = serializedOptions.get().as(FlinkPipelineOptions.class).isStreaming(); this.windowingStrategy = windowingStrategy; @@ -466,7 +469,12 @@ public void initializeState(StateInitializationContext context) throws Exception serializedOptions); sideInputHandler = new SideInputHandler(sideInputs, sideInputStateInternals); - sideInputReader = sideInputHandler; + sideInputReader = + createSideInputReader( + cacheableSideInputs, + getContainingTask().getEnvironment().getJobID(), + getContainingTask().getEnvironment().getTaskInfo().getAttemptNumber(), + sideInputHandler); Stream> pushedBack = pushedBackElementsHandler.getElements(); long min = @@ -790,6 +798,27 @@ protected void addSideInputValue(StreamRecord streamRecord) { PCollectionView sideInput = sideInputTagMapping.get(streamRecord.getValue().getUnionTag()); sideInputHandler.addSideInputValue(sideInput, value); + // Invalidate only after the state write: a concurrent reader that re-caches between an + // earlier invalidation and the write would pin the previous value with no later invalidation. + for (BoundedWindow window : value.getWindows()) { + SideInputCache.invalidate( + getContainingTask().getEnvironment().getJobID(), + getContainingTask().getEnvironment().getTaskInfo().getAttemptNumber(), + sideInput, + window); + } + } + + @VisibleForTesting + static SideInputReader createSideInputReader( + Collection> cacheableViews, + JobID jobId, + int attemptNumber, + SideInputReader delegate) { + if (!cacheableViews.isEmpty()) { + return CachedSideInputReader.of(jobId, attemptNumber, delegate, cacheableViews); + } + return delegate; } @Override diff --git a/runners/flink/2.2/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java b/runners/flink/2.2/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java index 6cebadc49d5c..06f3a39beab5 100644 --- a/runners/flink/2.2/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java +++ b/runners/flink/2.2/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java @@ -111,7 +111,6 @@ public void testDefaults() { assertThat(options.getAllowNonRestoredState(), is(false)); assertThat(options.getDisableMetrics(), is(false)); assertThat(options.getFasterCopy(), is(false)); - assertThat(options.isStreaming(), is(false)); assertThat(options.getMaxBundleSize(), is(5000L)); assertThat(options.getMaxBundleTimeMills(), is(10000L));