(beam) branch master updated: Handle null split initialization results in Flink lazy source enumerator (#40417)

[chamikara]

This is an automated email from the ASF dual-hosted git repository.

chamikaramj pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 6848f8454da Handle null split initialization results in Flink lazy 
source enumerator (#40417)
6848f8454da is described below

commit 6848f8454daa33e11a002c753d049e8a7384e65f
Author: ddebowczyk92 <ddebowczy...@users.noreply.github.com>
AuthorDate: Fri Oct 23 02:35:57 2026 +0200

    Handle null split initialization results in Flink lazy source enumerator 
(#40417)
---
 .../io/source/LazyFlinkSourceSplitEnumerator.java  |  9 ++++++-
 .../io/source/FlinkSourceSplitEnumeratorTest.java  | 28 ++++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git 
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java
 
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java
index 95d789178c0..8ec469abbd2 100644
--- 
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java
+++ 
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java
@@ -102,7 +102,14 @@ public final class LazyFlinkSourceSplitEnumerator<T>
         this::splitBeamSource,
         (sourceSplits, error) -> {
           if (error != null) {
-            throw new RuntimeException("Failed to start source enumerator.", 
error);
+            throw new RuntimeException(
+                "Failed to initialize splits for bounded source " + 
beamSource.getClass().getName(),
+                error);
+          }
+          if (sourceSplits == null) {
+            throw new IllegalStateException(
+                "Split initialization returned null without an error for 
bounded source "
+                    + beamSource.getClass().getName());
           }
           pendingSplits.addAll(sourceSplits);
           splitsInitialized = true;
diff --git 
a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java
 
b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java
index c4d220d2213..17a14c317bf 100644
--- 
a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java
+++ 
b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java
@@ -30,6 +30,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiConsumer;
 import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
 import org.apache.beam.runners.flink.FlinkPipelineOptions;
 import org.apache.beam.runners.flink.translation.utils.SerdeUtils;
@@ -43,9 +44,11 @@ import org.apache.beam.sdk.io.Source;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.values.KV;
 import org.apache.flink.api.connector.source.SplitEnumerator;
+import org.apache.flink.api.connector.source.SplitEnumeratorContext;
 import 
org.apache.flink.connector.testutils.source.reader.TestingSplitEnumeratorContext;
 import org.apache.flink.core.io.SimpleVersionedSerializer;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 /** Unit tests for the Flink source split enumerators. */
 public class FlinkSourceSplitEnumeratorTest {
@@ -54,6 +57,31 @@ public class FlinkSourceSplitEnumeratorTest {
   private static final int SOURCE_PARALLELISM = 2;
   private static final int REQUESTED_SPLITS = 4;
 
+  @Test
+  @SuppressWarnings("unchecked")
+  public void testLazyInitializationRejectsNullResultWithoutError() throws 
Exception {
+    SplitEnumeratorContext<FlinkSourceSplit<String>> context =
+        Mockito.mock(SplitEnumeratorContext.class);
+    Mockito.doAnswer(
+            invocation -> {
+              BiConsumer<ArrayList<FlinkSourceSplit<String>>, Throwable> 
callback =
+                  invocation.getArgument(1);
+              callback.accept(null, null);
+              return null;
+            })
+        .when(context)
+        .callAsync(Mockito.any(), Mockito.any());
+    try (LazyFlinkSourceSplitEnumerator<String> enumerator =
+        new LazyFlinkSourceSplitEnumerator<>(
+            context,
+            TestEstimatedSizeBoundedSource.create(1L, 1),
+            FlinkPipelineOptions.defaults(),
+            1)) {
+      IllegalStateException thrown = assertThrows(IllegalStateException.class, 
enumerator::start);
+      assertTrue(thrown.getMessage().contains("returned null without an 
error"));
+    }
+  }
+
   @Test
   public void testSelectsAssignmentModeFromEstimatedSizeAndConfiguration() 
throws Exception {
     assertEquals(0L, (long) 
FlinkPipelineOptions.defaults().getSourceStaticSplitThresholdMb());