From a79f23d9ea22b83f7ca46a1d8955140342d24f85 Mon Sep 17 00:00:00 2001 From: Paulius Kuzmickas Date: Thu, 13 Aug 2026 13:29:01 +0100 Subject: [PATCH 1/3] Cache Flink 2 batch side input materialization --- CHANGES.md | 1 + .../runners/flink/FlinkPipelineOptions.java | 10 + .../streaming/CachedSideInputReader.java | 55 +++++ .../wrappers/streaming/DoFnOperator.java | 25 ++- .../wrappers/streaming/SideInputCache.java | 113 +++++++++++ .../flink/FlinkPipelineOptionsTest.java | 1 + .../FlinkCachedSideInputReaderTest.java | 192 ++++++++++++++++++ .../wrappers/streaming/DoFnOperator.java | 25 ++- .../flink/FlinkPipelineOptionsTest.java | 1 + .../runners/flink/FlinkPipelineOptions.java | 10 + .../flink/FlinkPipelineOptionsTest.java | 1 + .../flink_java_pipeline_options.html | 5 + .../flink_python_pipeline_options.html | 5 + 13 files changed, 442 insertions(+), 2 deletions(-) create mode 100644 runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/CachedSideInputReader.java create mode 100644 runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SideInputCache.java create mode 100644 runners/flink/2.0/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkCachedSideInputReaderTest.java diff --git a/CHANGES.md b/CHANGES.md index 81f7d493677b..ec06b0e64082 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -103,6 +103,7 @@ ## New Features / Improvements +* Added opt-in caching of materialized side-input views for Flink DataStream batch execution with `--cacheSideInputMaterialization=true` (Java) ([#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/FlinkPipelineOptions.java b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java index 3fee130d58ee..084eaa791791 100644 --- a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java +++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java @@ -350,6 +350,16 @@ public Long create(PipelineOptions options) { void setFasterCopy(Boolean fasterCopy); + @Description( + "Batch/DataStream mode only (Flink 2.x): cache materialized side-input views per " + + "(view, window) in a process-wide cache, instead of re-applying the ViewFn against " + + "operator state on every access. Restores the per-TaskManager broadcast-variable " + + "caching of the legacy DataSet runner. No effect in streaming mode.") + @Default.Boolean(false) + Boolean getCacheSideInputMaterialization(); + + void setCacheSideInputMaterialization(Boolean cacheSideInputMaterialization); + @Description( "Directory containing Flink YAML configuration files. " + "These properties will be set to all jobs submitted to Flink and take precedence " 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..823c30331c2c --- /dev/null +++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/CachedSideInputReader.java @@ -0,0 +1,55 @@ +/* + * 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 org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.flink.api.common.JobID; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** {@link SideInputReader} that caches materialized views within a TaskManager JVM. */ +public final class CachedSideInputReader implements SideInputReader { + + public static CachedSideInputReader of(JobID jobId, SideInputReader delegate) { + return new CachedSideInputReader(jobId, delegate); + } + + private final JobID jobId; + private final SideInputReader delegate; + + private CachedSideInputReader(JobID jobId, SideInputReader delegate) { + this.jobId = jobId; + this.delegate = delegate; + } + + @Override + public @Nullable T get(PCollectionView view, BoundedWindow window) { + return SideInputCache.getOrMaterialize(jobId, 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..275e323429bf 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; @@ -466,7 +467,12 @@ public void initializeState(StateInitializationContext context) throws Exception serializedOptions); sideInputHandler = new SideInputHandler(sideInputs, sideInputStateInternals); - sideInputReader = sideInputHandler; + sideInputReader = + createSideInputReader( + isStreaming, + serializedOptions.get().as(FlinkPipelineOptions.class), + getContainingTask().getEnvironment().getJobID(), + sideInputHandler); Stream> pushedBack = pushedBackElementsHandler.getElements(); long min = @@ -630,6 +636,9 @@ private void earlyBindStateIfNeeded() throws IllegalArgumentException, IllegalAc } void cleanUp() throws Exception { + if (sideInputReader instanceof CachedSideInputReader) { + SideInputCache.invalidateAll(getContainingTask().getEnvironment().getJobID()); + } Optional.ofNullable(flinkMetricContainer) .ifPresent(FlinkMetricContainer::registerMetricsForPipelineResult); Optional.ofNullable(checkFinishBundleTimer).ifPresent(timer -> timer.cancel(true)); @@ -790,6 +799,20 @@ 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(), sideInput, window); + } + } + + @VisibleForTesting + static SideInputReader createSideInputReader( + boolean isStreaming, FlinkPipelineOptions options, JobID jobId, SideInputReader delegate) { + if (!isStreaming && options.getCacheSideInputMaterialization()) { + return CachedSideInputReader.of(jobId, delegate); + } + 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..cf326ca55464 --- /dev/null +++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SideInputCache.java @@ -0,0 +1,113 @@ +/* + * 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. + private static final Cache, Value> MATERIALIZED_SIDE_INPUTS = + CacheBuilder.newBuilder().expireAfterAccess(5, TimeUnit.MINUTES).softValues().build(); + + private SideInputCache() {} + + static @Nullable T getOrMaterialize( + JobID jobId, + 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, 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, PCollectionView view, BoundedWindow window) { + MATERIALIZED_SIDE_INPUTS.invalidate(new Key<>(jobId, view, window)); + } + + static void invalidateAll(JobID jobId) { + MATERIALIZED_SIDE_INPUTS.asMap().keySet().removeIf(key -> jobId.equals(key.jobId)); + } + + private static final class Key { + private final JobID jobId; + private final PCollectionView view; + private final BoundedWindow window; + + private Key(JobID jobId, PCollectionView view, BoundedWindow window) { + this.jobId = jobId; + 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) + && Objects.equals(view, other.view) + && Objects.equals(window, other.window); + } + + @Override + public int hashCode() { + return Objects.hash(jobId, 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..73b91143d021 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,6 +97,7 @@ public void testDefaults() { assertThat(options.getAllowNonRestoredState(), is(false)); assertThat(options.getDisableMetrics(), is(false)); assertThat(options.getFasterCopy(), is(false)); + assertThat(options.getCacheSideInputMaterialization(), is(false)); assertThat(options.isStreaming(), is(false)); assertThat(options.getMaxBundleSize(), is(5000L)); 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..8d6ba2051e70 --- /dev/null +++ b/runners/flink/2.0/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkCachedSideInputReaderTest.java @@ -0,0 +1,192 @@ +/* + * 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.mock; + +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.runners.flink.FlinkPipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.transforms.windowing.IntervalWindow; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.flink.api.common.JobID; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; +import org.junit.Test; + +/** Tests for cached materialization of Flink side-input views. */ +public class FlinkCachedSideInputReaderTest { + + @Test + public void repeatedGetMaterializesOnce() { + JobID jobId = new JobID(); + PCollectionView view = view(); + CountingSideInputReader delegate = new CountingSideInputReader("value"); + SideInputReader reader = CachedSideInputReader.of(jobId, delegate); + + assertThat(reader.get(view, GlobalWindow.INSTANCE), is("value")); + assertThat(reader.get(view, GlobalWindow.INSTANCE), is("value")); + assertThat(delegate.getCount(), is(1)); + } + + @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(); + + CachedSideInputReader.of(firstJob, delegate).get(firstView, firstWindow); + CachedSideInputReader.of(firstJob, delegate).get(secondView, firstWindow); + CachedSideInputReader.of(firstJob, delegate).get(firstView, secondWindow); + CachedSideInputReader.of(new JobID(), delegate).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, delegate); + + reader.get(view, GlobalWindow.INSTANCE); + SideInputCache.invalidate(jobId, 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, delegate); + + assertThat(reader.get(view, GlobalWindow.INSTANCE), nullValue()); + assertThat(reader.get(view, GlobalWindow.INSTANCE), nullValue()); + assertThat(delegate.getCount(), is(1)); + } + + @Test + public void optionWrapsOnlyBatchReaderWhenEnabled() { + JobID jobId = new JobID(); + SideInputReader delegate = new CountingSideInputReader("value"); + FlinkPipelineOptions options = PipelineOptionsFactory.as(FlinkPipelineOptions.class); + + assertThat(DoFnOperator.createSideInputReader(false, options, jobId, delegate), is(delegate)); + + options.setCacheSideInputMaterialization(true); + assertThat( + DoFnOperator.createSideInputReader(false, options, jobId, delegate), + instanceOf(CachedSideInputReader.class)); + assertThat(DoFnOperator.createSideInputReader(true, options, jobId, delegate), is(delegate)); + } + + @Test + public void invalidateAllRemovesOnlyEntriesOfJob() { + JobID firstJob = new JobID(); + JobID secondJob = new JobID(); + PCollectionView view = view(); + CountingSideInputReader delegate = new CountingSideInputReader("value"); + CachedSideInputReader.of(firstJob, delegate).get(view, GlobalWindow.INSTANCE); + CachedSideInputReader.of(secondJob, delegate).get(view, GlobalWindow.INSTANCE); + + SideInputCache.invalidateAll(firstJob); + + CachedSideInputReader.of(secondJob, delegate).get(view, GlobalWindow.INSTANCE); + assertThat(delegate.getCount(), is(2)); + CachedSideInputReader.of(firstJob, delegate).get(view, GlobalWindow.INSTANCE); + assertThat(delegate.getCount(), is(3)); + } + + @Test + public void materializationExceptionPropagatesUnwrapped() { + SideInputReader reader = + CachedSideInputReader.of( + new JobID(), + 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; + } + }); + + IllegalStateException exception = + assertThrows(IllegalStateException.class, () -> reader.get(view(), GlobalWindow.INSTANCE)); + assertThat(exception.getMessage(), is("materialization failed")); + } + + @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..14fc794eaa1a 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; @@ -466,7 +467,12 @@ public void initializeState(StateInitializationContext context) throws Exception serializedOptions); sideInputHandler = new SideInputHandler(sideInputs, sideInputStateInternals); - sideInputReader = sideInputHandler; + sideInputReader = + createSideInputReader( + isStreaming, + serializedOptions.get().as(FlinkPipelineOptions.class), + getContainingTask().getEnvironment().getJobID(), + sideInputHandler); Stream> pushedBack = pushedBackElementsHandler.getElements(); long min = @@ -630,6 +636,9 @@ private void earlyBindStateIfNeeded() throws IllegalArgumentException, IllegalAc } void cleanUp() throws Exception { + if (sideInputReader instanceof CachedSideInputReader) { + SideInputCache.invalidateAll(getContainingTask().getEnvironment().getJobID()); + } Optional.ofNullable(flinkMetricContainer) .ifPresent(FlinkMetricContainer::registerMetricsForPipelineResult); Optional.ofNullable(checkFinishBundleTimer).ifPresent(timer -> timer.cancel(true)); @@ -790,6 +799,20 @@ 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(), sideInput, window); + } + } + + @VisibleForTesting + static SideInputReader createSideInputReader( + boolean isStreaming, FlinkPipelineOptions options, JobID jobId, SideInputReader delegate) { + if (!isStreaming && options.getCacheSideInputMaterialization()) { + return CachedSideInputReader.of(jobId, delegate); + } + 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..70b4dc51b0a8 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,6 +111,7 @@ public void testDefaults() { assertThat(options.getAllowNonRestoredState(), is(false)); assertThat(options.getDisableMetrics(), is(false)); assertThat(options.getFasterCopy(), is(false)); + assertThat(options.getCacheSideInputMaterialization(), is(false)); assertThat(options.isStreaming(), is(false)); assertThat(options.getMaxBundleSize(), is(5000L)); diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java index f0724b4d031f..c58aff4ef89a 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java @@ -365,6 +365,16 @@ public Long create(PipelineOptions options) { void setFasterCopy(Boolean fasterCopy); + @Description( + "Batch/DataStream mode only (Flink 2.x): cache materialized side-input views per " + + "(view, window) in a process-wide cache, instead of re-applying the ViewFn against " + + "operator state on every access. Restores the per-TaskManager broadcast-variable " + + "caching of the legacy DataSet runner. No effect in streaming mode.") + @Default.Boolean(false) + Boolean getCacheSideInputMaterialization(); + + void setCacheSideInputMaterialization(Boolean cacheSideInputMaterialization); + @Description( "Directory containing Flink YAML configuration files. " + "These properties will be set to all jobs submitted to Flink and take precedence " diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java index dc07caf5bfd7..54e2bf707d86 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java @@ -98,6 +98,7 @@ public void testDefaults() { assertThat(options.getAllowNonRestoredState(), is(false)); assertThat(options.getDisableMetrics(), is(false)); assertThat(options.getFasterCopy(), is(false)); + assertThat(options.getCacheSideInputMaterialization(), is(false)); assertThat(options.isStreaming(), is(false)); assertThat(options.getMaxBundleSize(), is(5000L)); diff --git a/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html b/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html index 34d6c5243776..cd0de955741a 100644 --- a/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html +++ b/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html @@ -37,6 +37,11 @@ The interval in milliseconds for automatic watermark emission. + + cacheSideInputMaterialization + Batch/DataStream mode only (Flink 2.x): cache materialized side-input views per (view, window) in a process-wide cache, instead of re-applying the ViewFn against operator state on every access. Restores the per-TaskManager broadcast-variable caching of the legacy DataSet runner. No effect in streaming mode. + Default: false + checkpointTimeoutMillis The maximum time in milliseconds that a checkpoint may take before being discarded. diff --git a/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html b/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html index e3fe24216a54..9c5b5fe22343 100644 --- a/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html +++ b/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html @@ -37,6 +37,11 @@ The interval in milliseconds for automatic watermark emission. + + cache_side_input_materialization + Batch/DataStream mode only (Flink 2.x): cache materialized side-input views per (view, window) in a process-wide cache, instead of re-applying the ViewFn against operator state on every access. Restores the per-TaskManager broadcast-variable caching of the legacy DataSet runner. No effect in streaming mode. + Default: false + checkpoint_timeout_millis The maximum time in milliseconds that a checkpoint may take before being discarded. From 09340dccf7adba4564464bf31db579a7ee9e84fb Mon Sep 17 00:00:00 2001 From: Paulius Kuzmickas Date: Tue, 25 Aug 2026 16:08:41 +0100 Subject: [PATCH 2/3] Scope side input cache option to Flink 2 --- CHANGES.md | 2 +- .../beam/runners/flink/FlinkPipelineOptions.java | 4 ++-- .../beam/runners/flink/FlinkPipelineOptions.java | 10 ---------- .../beam/runners/flink/FlinkPipelineOptionsTest.java | 1 - .../shortcodes/flink_java_pipeline_options.html | 2 +- .../shortcodes/flink_python_pipeline_options.html | 2 +- 6 files changed, 5 insertions(+), 16 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index ec06b0e64082..219b1f8d5277 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -103,7 +103,7 @@ ## New Features / Improvements -* Added opt-in caching of materialized side-input views for Flink DataStream batch execution with `--cacheSideInputMaterialization=true` (Java) ([#39866](https://github.com/apache/beam/issues/39866)). +* Added opt-in caching of materialized side-input views for classic Java Flink DataStream batch execution with `--cacheSideInputMaterialization=true` ([#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/FlinkPipelineOptions.java b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java index 084eaa791791..3b4b6ed1c26f 100644 --- a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java +++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java @@ -351,10 +351,10 @@ public Long create(PipelineOptions options) { void setFasterCopy(Boolean fasterCopy); @Description( - "Batch/DataStream mode only (Flink 2.x): cache materialized side-input views per " + "Classic Java batch runner only (Flink 2.x): cache materialized side-input views per " + "(view, window) in a process-wide cache, instead of re-applying the ViewFn against " + "operator state on every access. Restores the per-TaskManager broadcast-variable " - + "caching of the legacy DataSet runner. No effect in streaming mode.") + + "caching of the legacy DataSet runner. No effect in portable or streaming mode.") @Default.Boolean(false) Boolean getCacheSideInputMaterialization(); diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java index c58aff4ef89a..f0724b4d031f 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java @@ -365,16 +365,6 @@ public Long create(PipelineOptions options) { void setFasterCopy(Boolean fasterCopy); - @Description( - "Batch/DataStream mode only (Flink 2.x): cache materialized side-input views per " - + "(view, window) in a process-wide cache, instead of re-applying the ViewFn against " - + "operator state on every access. Restores the per-TaskManager broadcast-variable " - + "caching of the legacy DataSet runner. No effect in streaming mode.") - @Default.Boolean(false) - Boolean getCacheSideInputMaterialization(); - - void setCacheSideInputMaterialization(Boolean cacheSideInputMaterialization); - @Description( "Directory containing Flink YAML configuration files. " + "These properties will be set to all jobs submitted to Flink and take precedence " diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java index 54e2bf707d86..dc07caf5bfd7 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java @@ -98,7 +98,6 @@ public void testDefaults() { assertThat(options.getAllowNonRestoredState(), is(false)); assertThat(options.getDisableMetrics(), is(false)); assertThat(options.getFasterCopy(), is(false)); - assertThat(options.getCacheSideInputMaterialization(), is(false)); assertThat(options.isStreaming(), is(false)); assertThat(options.getMaxBundleSize(), is(5000L)); diff --git a/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html b/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html index cd0de955741a..c8e9ce4b5cfb 100644 --- a/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html +++ b/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html @@ -39,7 +39,7 @@ cacheSideInputMaterialization - Batch/DataStream mode only (Flink 2.x): cache materialized side-input views per (view, window) in a process-wide cache, instead of re-applying the ViewFn against operator state on every access. Restores the per-TaskManager broadcast-variable caching of the legacy DataSet runner. No effect in streaming mode. + Classic Java batch runner only (Flink 2.x): cache materialized side-input views per (view, window) in a process-wide cache, instead of re-applying the ViewFn against operator state on every access. Restores the per-TaskManager broadcast-variable caching of the legacy DataSet runner. No effect in portable or streaming mode. Default: false diff --git a/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html b/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html index 9c5b5fe22343..9f9a395c6a8e 100644 --- a/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html +++ b/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html @@ -39,7 +39,7 @@ cache_side_input_materialization - Batch/DataStream mode only (Flink 2.x): cache materialized side-input views per (view, window) in a process-wide cache, instead of re-applying the ViewFn against operator state on every access. Restores the per-TaskManager broadcast-variable caching of the legacy DataSet runner. No effect in streaming mode. + Classic Java batch runner only (Flink 2.x): cache materialized side-input views per (view, window) in a process-wide cache, instead of re-applying the ViewFn against operator state on every access. Restores the per-TaskManager broadcast-variable caching of the legacy DataSet runner. No effect in portable or streaming mode. Default: false From 7d1fbe19929464bd4446d4c28792be6c394cd8d6 Mon Sep 17 00:00:00 2001 From: Paulius Kuzmickas Date: Mon, 31 Aug 2026 16:56:10 +0100 Subject: [PATCH 3/3] Cache safe bounded Flink side inputs automatically --- CHANGES.md | 2 +- .../runners/flink/FlinkPipelineOptions.java | 10 -- .../streaming/CachedSideInputReader.java | 48 +++++- .../wrappers/streaming/DoFnOperator.java | 24 +-- .../wrappers/streaming/SideInputCache.java | 22 +-- .../flink/FlinkPipelineOptionsTest.java | 2 - .../FlinkCachedSideInputReaderTest.java | 147 ++++++++++++++---- .../wrappers/streaming/DoFnOperator.java | 24 +-- .../flink/FlinkPipelineOptionsTest.java | 2 - .../flink_java_pipeline_options.html | 5 - .../flink_python_pipeline_options.html | 5 - 11 files changed, 206 insertions(+), 85 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 219b1f8d5277..4eb1db840d6f 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -103,7 +103,7 @@ ## New Features / Improvements -* Added opt-in caching of materialized side-input views for classic Java Flink DataStream batch execution with `--cacheSideInputMaterialization=true` ([#39866](https://github.com/apache/beam/issues/39866)). +* 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/FlinkPipelineOptions.java b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java index 3b4b6ed1c26f..3fee130d58ee 100644 --- a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java +++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java @@ -350,16 +350,6 @@ public Long create(PipelineOptions options) { void setFasterCopy(Boolean fasterCopy); - @Description( - "Classic Java batch runner only (Flink 2.x): cache materialized side-input views per " - + "(view, window) in a process-wide cache, instead of re-applying the ViewFn against " - + "operator state on every access. Restores the per-TaskManager broadcast-variable " - + "caching of the legacy DataSet runner. No effect in portable or streaming mode.") - @Default.Boolean(false) - Boolean getCacheSideInputMaterialization(); - - void setCacheSideInputMaterialization(Boolean cacheSideInputMaterialization); - @Description( "Directory containing Flink YAML configuration files. " + "These properties will be set to all jobs submitted to Flink and take precedence " 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 index 823c30331c2c..ee0d9df53b57 100644 --- 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 @@ -17,30 +17,68 @@ */ 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 materialized views within a TaskManager JVM. */ +/** {@link SideInputReader} that caches single-pane materialized views within a TaskManager JVM. */ public final class CachedSideInputReader implements SideInputReader { - public static CachedSideInputReader of(JobID jobId, SideInputReader delegate) { - return new CachedSideInputReader(jobId, delegate); + 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, SideInputReader delegate) { + 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) { - return SideInputCache.getOrMaterialize(jobId, view, window, () -> delegate.get(view, window)); + if (!cacheableViews.contains(view)) { + return delegate.get(view, window); + } + return SideInputCache.getOrMaterialize( + jobId, attemptNumber, view, window, () -> delegate.get(view, window)); } @Override 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 275e323429bf..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 @@ -163,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; @@ -298,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; @@ -469,9 +471,9 @@ public void initializeState(StateInitializationContext context) throws Exception sideInputHandler = new SideInputHandler(sideInputs, sideInputStateInternals); sideInputReader = createSideInputReader( - isStreaming, - serializedOptions.get().as(FlinkPipelineOptions.class), + cacheableSideInputs, getContainingTask().getEnvironment().getJobID(), + getContainingTask().getEnvironment().getTaskInfo().getAttemptNumber(), sideInputHandler); Stream> pushedBack = pushedBackElementsHandler.getElements(); @@ -636,9 +638,6 @@ private void earlyBindStateIfNeeded() throws IllegalArgumentException, IllegalAc } void cleanUp() throws Exception { - if (sideInputReader instanceof CachedSideInputReader) { - SideInputCache.invalidateAll(getContainingTask().getEnvironment().getJobID()); - } Optional.ofNullable(flinkMetricContainer) .ifPresent(FlinkMetricContainer::registerMetricsForPipelineResult); Optional.ofNullable(checkFinishBundleTimer).ifPresent(timer -> timer.cancel(true)); @@ -802,15 +801,22 @@ protected void addSideInputValue(StreamRecord streamRecord) { // 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(), sideInput, window); + SideInputCache.invalidate( + getContainingTask().getEnvironment().getJobID(), + getContainingTask().getEnvironment().getTaskInfo().getAttemptNumber(), + sideInput, + window); } } @VisibleForTesting static SideInputReader createSideInputReader( - boolean isStreaming, FlinkPipelineOptions options, JobID jobId, SideInputReader delegate) { - if (!isStreaming && options.getCacheSideInputMaterialization()) { - return CachedSideInputReader.of(jobId, delegate); + Collection> cacheableViews, + JobID jobId, + int attemptNumber, + SideInputReader delegate) { + if (!cacheableViews.isEmpty()) { + return CachedSideInputReader.of(jobId, attemptNumber, delegate, cacheableViews); } return delegate; } 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 index cf326ca55464..ce49d8a5de99 100644 --- 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 @@ -35,6 +35,9 @@ 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(); @@ -42,6 +45,7 @@ private SideInputCache() {} static @Nullable T getOrMaterialize( JobID jobId, + int attemptNumber, PCollectionView view, BoundedWindow window, Supplier<@Nullable T> materializer) { @@ -50,7 +54,7 @@ private SideInputCache() {} (Cache, Value>) (Cache) MATERIALIZED_SIDE_INPUTS; try { return cache - .get(new Key<>(jobId, view, window), () -> new Value<>(materializer.get())) + .get(new Key<>(jobId, attemptNumber, view, window), () -> new Value<>(materializer.get())) .getValue(); } catch (ExecutionException | UncheckedExecutionException e) { Throwable cause = e.getCause() != null ? e.getCause() : e; @@ -59,21 +63,20 @@ private SideInputCache() {} } } - static void invalidate(JobID jobId, PCollectionView view, BoundedWindow window) { - MATERIALIZED_SIDE_INPUTS.invalidate(new Key<>(jobId, view, window)); - } - - static void invalidateAll(JobID jobId) { - MATERIALIZED_SIDE_INPUTS.asMap().keySet().removeIf(key -> jobId.equals(key.jobId)); + 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, PCollectionView view, BoundedWindow window) { + private Key(JobID jobId, int attemptNumber, PCollectionView view, BoundedWindow window) { this.jobId = jobId; + this.attemptNumber = attemptNumber; this.view = view; this.window = window; } @@ -88,13 +91,14 @@ public boolean equals(@Nullable Object object) { } 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, view, window); + return Objects.hash(jobId, attemptNumber, view, window); } } 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 73b91143d021..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,8 +97,6 @@ public void testDefaults() { assertThat(options.getAllowNonRestoredState(), is(false)); assertThat(options.getDisableMetrics(), is(false)); assertThat(options.getFasterCopy(), is(false)); - assertThat(options.getCacheSideInputMaterialization(), 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 index 8d6ba2051e70..c71dcd295449 100644 --- 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 @@ -22,36 +22,78 @@ 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.runners.flink.FlinkPipelineOptions; -import org.apache.beam.sdk.options.PipelineOptionsFactory; 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, delegate); + 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(); @@ -62,10 +104,15 @@ public void keyIncludesViewWindowAndJob() { CountingSideInputReader delegate = new CountingSideInputReader("value"); JobID firstJob = new JobID(); - CachedSideInputReader.of(firstJob, delegate).get(firstView, firstWindow); - CachedSideInputReader.of(firstJob, delegate).get(secondView, firstWindow); - CachedSideInputReader.of(firstJob, delegate).get(firstView, secondWindow); - CachedSideInputReader.of(new JobID(), delegate).get(firstView, firstWindow); + 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)); } @@ -75,10 +122,11 @@ public void invalidateRematerializesValue() { JobID jobId = new JobID(); PCollectionView view = view(); CountingSideInputReader delegate = new CountingSideInputReader("value"); - SideInputReader reader = CachedSideInputReader.of(jobId, delegate); + SideInputReader reader = + CachedSideInputReader.of(jobId, INITIAL_ATTEMPT, delegate, Collections.singleton(view)); reader.get(view, GlobalWindow.INSTANCE); - SideInputCache.invalidate(jobId, view, GlobalWindow.INSTANCE); + SideInputCache.invalidate(jobId, INITIAL_ATTEMPT, view, GlobalWindow.INSTANCE); reader.get(view, GlobalWindow.INSTANCE); assertThat(delegate.getCount(), is(2)); @@ -89,7 +137,8 @@ public void cachesNull() { JobID jobId = new JobID(); PCollectionView view = view(); CountingSideInputReader delegate = new CountingSideInputReader(null); - SideInputReader reader = CachedSideInputReader.of(jobId, delegate); + 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()); @@ -97,42 +146,65 @@ public void cachesNull() { } @Test - public void optionWrapsOnlyBatchReaderWhenEnabled() { + public void automaticallyWrapsReaderWithCacheableViews() { JobID jobId = new JobID(); SideInputReader delegate = new CountingSideInputReader("value"); - FlinkPipelineOptions options = PipelineOptionsFactory.as(FlinkPipelineOptions.class); + PCollectionView view = cacheableView(); + Collection> cacheableViews = + CachedSideInputReader.cacheableViews(Collections.singleton(view)); - assertThat(DoFnOperator.createSideInputReader(false, options, jobId, delegate), is(delegate)); + assertThat( + DoFnOperator.createSideInputReader( + Collections.emptyList(), jobId, INITIAL_ATTEMPT, delegate), + is(delegate)); - options.setCacheSideInputMaterialization(true); assertThat( - DoFnOperator.createSideInputReader(false, options, jobId, delegate), + DoFnOperator.createSideInputReader(cacheableViews, jobId, INITIAL_ATTEMPT, delegate), instanceOf(CachedSideInputReader.class)); - assertThat(DoFnOperator.createSideInputReader(true, options, jobId, delegate), is(delegate)); } @Test - public void invalidateAllRemovesOnlyEntriesOfJob() { - JobID firstJob = new JobID(); - JobID secondJob = new JobID(); - PCollectionView view = view(); + 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"); - CachedSideInputReader.of(firstJob, delegate).get(view, GlobalWindow.INSTANCE); - CachedSideInputReader.of(secondJob, delegate).get(view, GlobalWindow.INSTANCE); + SideInputReader reader = + CachedSideInputReader.of( + new JobID(), INITIAL_ATTEMPT, delegate, Collections.singleton(cacheableView)); - SideInputCache.invalidateAll(firstJob); + reader.get(cacheableView, GlobalWindow.INSTANCE); + reader.get(cacheableView, GlobalWindow.INSTANCE); + reader.get(nonCacheableView, GlobalWindow.INSTANCE); + reader.get(nonCacheableView, GlobalWindow.INSTANCE); - CachedSideInputReader.of(secondJob, delegate).get(view, GlobalWindow.INSTANCE); - assertThat(delegate.getCount(), is(2)); - CachedSideInputReader.of(firstJob, delegate).get(view, 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) { @@ -148,13 +220,32 @@ public boolean contains(PCollectionView view) { public boolean isEmpty() { return false; } - }); + }, + Collections.singleton(view)); IllegalStateException exception = - assertThrows(IllegalStateException.class, () -> reader.get(view(), GlobalWindow.INSTANCE)); + 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); 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 14fc794eaa1a..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 @@ -163,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; @@ -298,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; @@ -469,9 +471,9 @@ public void initializeState(StateInitializationContext context) throws Exception sideInputHandler = new SideInputHandler(sideInputs, sideInputStateInternals); sideInputReader = createSideInputReader( - isStreaming, - serializedOptions.get().as(FlinkPipelineOptions.class), + cacheableSideInputs, getContainingTask().getEnvironment().getJobID(), + getContainingTask().getEnvironment().getTaskInfo().getAttemptNumber(), sideInputHandler); Stream> pushedBack = pushedBackElementsHandler.getElements(); @@ -636,9 +638,6 @@ private void earlyBindStateIfNeeded() throws IllegalArgumentException, IllegalAc } void cleanUp() throws Exception { - if (sideInputReader instanceof CachedSideInputReader) { - SideInputCache.invalidateAll(getContainingTask().getEnvironment().getJobID()); - } Optional.ofNullable(flinkMetricContainer) .ifPresent(FlinkMetricContainer::registerMetricsForPipelineResult); Optional.ofNullable(checkFinishBundleTimer).ifPresent(timer -> timer.cancel(true)); @@ -802,15 +801,22 @@ protected void addSideInputValue(StreamRecord streamRecord) { // 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(), sideInput, window); + SideInputCache.invalidate( + getContainingTask().getEnvironment().getJobID(), + getContainingTask().getEnvironment().getTaskInfo().getAttemptNumber(), + sideInput, + window); } } @VisibleForTesting static SideInputReader createSideInputReader( - boolean isStreaming, FlinkPipelineOptions options, JobID jobId, SideInputReader delegate) { - if (!isStreaming && options.getCacheSideInputMaterialization()) { - return CachedSideInputReader.of(jobId, delegate); + Collection> cacheableViews, + JobID jobId, + int attemptNumber, + SideInputReader delegate) { + if (!cacheableViews.isEmpty()) { + return CachedSideInputReader.of(jobId, attemptNumber, delegate, cacheableViews); } return delegate; } 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 70b4dc51b0a8..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,8 +111,6 @@ public void testDefaults() { assertThat(options.getAllowNonRestoredState(), is(false)); assertThat(options.getDisableMetrics(), is(false)); assertThat(options.getFasterCopy(), is(false)); - assertThat(options.getCacheSideInputMaterialization(), is(false)); - assertThat(options.isStreaming(), is(false)); assertThat(options.getMaxBundleSize(), is(5000L)); assertThat(options.getMaxBundleTimeMills(), is(10000L)); diff --git a/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html b/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html index c8e9ce4b5cfb..34d6c5243776 100644 --- a/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html +++ b/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html @@ -37,11 +37,6 @@ The interval in milliseconds for automatic watermark emission. - - cacheSideInputMaterialization - Classic Java batch runner only (Flink 2.x): cache materialized side-input views per (view, window) in a process-wide cache, instead of re-applying the ViewFn against operator state on every access. Restores the per-TaskManager broadcast-variable caching of the legacy DataSet runner. No effect in portable or streaming mode. - Default: false - checkpointTimeoutMillis The maximum time in milliseconds that a checkpoint may take before being discarded. diff --git a/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html b/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html index 9f9a395c6a8e..e3fe24216a54 100644 --- a/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html +++ b/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html @@ -37,11 +37,6 @@ The interval in milliseconds for automatic watermark emission. - - cache_side_input_materialization - Classic Java batch runner only (Flink 2.x): cache materialized side-input views per (view, window) in a process-wide cache, instead of re-applying the ViewFn against operator state on every access. Restores the per-TaskManager broadcast-variable caching of the legacy DataSet runner. No effect in portable or streaming mode. - Default: false - checkpoint_timeout_millis The maximum time in milliseconds that a checkpoint may take before being discarded.