[Flink] Cache batch side input materialization in Flink 2

CHANGES.md

@@ -103,6 +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 `GroupIntoBatches` transform and the standard
   `beam:coder:sharded_key:v1` coder to the Go SDK, along with
   `beam.Coder.IsDeterministic`, `beam.PCollection.WindowingStrategy`,

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(
+      "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 "

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 <T> @Nullable T get(PCollectionView<T> view, BoundedWindow window) {
+    return SideInputCache.getOrMaterialize(jobId, view, window, () -> delegate.get(view, window));
+  }
+
+  @Override
+  public <T> boolean contains(PCollectionView<T> view) {
+    return delegate.contains(view);
+  }
+
+  @Override
+  public boolean isEmpty() {
+    return delegate.isEmpty();
+  }
+}

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<WindowedValue<InputT>> 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<RawUnionValue> 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

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<Key<?>, Value<?>> MATERIALIZED_SIDE_INPUTS =
+      CacheBuilder.newBuilder().expireAfterAccess(5, TimeUnit.MINUTES).softValues().build();
+
+  private SideInputCache() {}
+
+  static <T> @Nullable T getOrMaterialize(
+      JobID jobId,
+      PCollectionView<T> view,
+      BoundedWindow window,
+      Supplier<@Nullable T> materializer) {
+    @SuppressWarnings("unchecked")
+    Cache<Key<T>, Value<T>> cache =
+        (Cache<Key<T>, Value<T>>) (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<T> {
+    private final JobID jobId;
+    private final PCollectionView<T> view;
+    private final BoundedWindow window;
+
+    private Key(JobID jobId, PCollectionView<T> 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<T> {
+    private final @Nullable T value;
+
+    private Value(@Nullable T value) {
+      this.value = value;
+    }
+
+    private @Nullable T getValue() {
+      return value;
+    }
+  }
+}

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));