[Flink] Cache batch side input materialization in Flink 2

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`,

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<PCollectionView<?>> cacheableViews) {
+    return new CachedSideInputReader(jobId, attemptNumber, delegate, cacheableViews);
+  }
+
+  static Collection<PCollectionView<?>> cacheableViews(Collection<PCollectionView<?>> sideInputs) {
+    Collection<PCollectionView<?>> 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<PCollectionView<?>> cacheableViews;
+
+  private CachedSideInputReader(
+      JobID jobId,
+      int attemptNumber,
+      SideInputReader delegate,
+      Collection<PCollectionView<?>> cacheableViews) {
+    this.jobId = jobId;
+    this.attemptNumber = attemptNumber;
+    this.delegate = delegate;
+    this.cacheableViews = cacheableViews;
+  }
+
+  @Override
+  public <T> @Nullable T get(PCollectionView<T> 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 <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;
@@ -162,6 +163,7 @@ public class DoFnOperator<PreInputT, InputT, OutputT>
   protected final List<TupleTag<?>> additionalOutputTags;
 
   protected final Collection<PCollectionView<?>> sideInputs;
+  private final Collection<PCollectionView<?>> cacheableSideInputs;
   protected final Map<Integer, PCollectionView<?>> 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<WindowedValue<InputT>> pushedBack = pushedBackElementsHandler.getElements();
       long min =
@@ -790,6 +798,27 @@ 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(),
+          getContainingTask().getEnvironment().getTaskInfo().getAttemptNumber(),
+          sideInput,
+          window);
+    }
+  }
+
+  @VisibleForTesting
+  static SideInputReader createSideInputReader(
+      Collection<PCollectionView<?>> cacheableViews,
+      JobID jobId,
+      int attemptNumber,
+      SideInputReader delegate) {
+    if (!cacheableViews.isEmpty()) {
+      return CachedSideInputReader.of(jobId, attemptNumber, delegate, cacheableViews);
+    }
+    return delegate;
   }
 
   @Override

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<Key<?>, Value<?>> MATERIALIZED_SIDE_INPUTS =
+      CacheBuilder.newBuilder().expireAfterAccess(5, TimeUnit.MINUTES).softValues().build();
+
+  private SideInputCache() {}
+
+  static <T> @Nullable T getOrMaterialize(
+      JobID jobId,
+      int attemptNumber,
+      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, 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<T> {
+    private final JobID jobId;
+    private final int attemptNumber;
+    private final PCollectionView<T> view;
+    private final BoundedWindow window;
+
+    private Key(JobID jobId, int attemptNumber, PCollectionView<T> 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<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,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));