You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@hudi.apache.org by xu...@apache.org on 2025/12/09 23:47:06 UTC
(hudi) branch master updated: feat: add static split enumerator for Flink source (#17503)
This is an automated email from the ASF dual-hosted git repository.
xushiyan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hudi.git
The following commit(s) were added to refs/heads/master by this push:
new c48866904fcb feat: add static split enumerator for Flink source (#17503)
c48866904fcb is described below
commit c48866904fcbefeb266abfee480d6a7ae70ac043
Author: Peter Huang <hu...@gmail.com>
AuthorDate: Tue Dec 9 15:46:59 2025 -0800
feat: add static split enumerator for Flink source (#17503)
---
.../enumerator/AbstractHoodieSplitEnumerator.java | 188 +++++++++++++
.../enumerator/HoodieSplitEnumeratorState.java | 39 +++
.../enumerator/HoodieStaticSplitEnumerator.java | 40 +++
.../source/split/DefaultHoodieSplitProvider.java | 84 ++++++
.../hudi/source/split/HoodieSplitProvider.java | 122 ++++++++
.../TestHoodieStaticSplitEnumerator.java | 311 +++++++++++++++++++++
.../split/TestDefaultHoodieSplitProvider.java | 188 +++++++++++++
7 files changed, 972 insertions(+)
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/enumerator/AbstractHoodieSplitEnumerator.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/enumerator/AbstractHoodieSplitEnumerator.java
new file mode 100644
index 000000000000..331ee0d7e22d
--- /dev/null
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/enumerator/AbstractHoodieSplitEnumerator.java
@@ -0,0 +1,188 @@
+/*
+ * 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.hudi.source.enumerator;
+
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.source.split.HoodieSourceSplit;
+import org.apache.hudi.source.split.HoodieSplitProvider;
+import org.apache.hudi.source.split.SplitRequestEvent;
+
+import org.apache.flink.api.connector.source.SourceEvent;
+import org.apache.flink.api.connector.source.SplitEnumerator;
+import org.apache.flink.api.connector.source.SplitEnumeratorContext;
+import org.apache.flink.api.connector.source.SupportsHandleExecutionAttemptSourceEvent;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * Abstract class for Hoodie SourceSplit enumerator.
+ */
+abstract class AbstractHoodieSplitEnumerator
+ implements SplitEnumerator<HoodieSourceSplit, HoodieSplitEnumeratorState>,
+ SupportsHandleExecutionAttemptSourceEvent {
+ private static final Logger LOG = LoggerFactory.getLogger(AbstractHoodieSplitEnumerator.class);
+
+ // Split provider that provide Hoodie split for split request
+ protected final HoodieSplitProvider splitProvider;
+ // Context of split enumerator
+ private final SplitEnumeratorContext<HoodieSourceSplit> enumeratorContext;
+ // Registered reader task id to request host home map
+ private final Map<Integer, String> readersAwaitingSplit;
+ // Completable future for signal of new split readiness
+ private final AtomicReference<CompletableFuture<Void>> availableFuture;
+
+ AbstractHoodieSplitEnumerator(
+ SplitEnumeratorContext<HoodieSourceSplit> enumeratorContext,
+ HoodieSplitProvider splitProvider) {
+ this.enumeratorContext = enumeratorContext;
+ // Use LinkedHashMap to make sure split request are handled in request order
+ this.readersAwaitingSplit = new LinkedHashMap<>();
+ this.availableFuture = new AtomicReference<>();
+ this.splitProvider = splitProvider;
+ this.enumeratorContext
+ .metricGroup()
+ // This number may not capture the entire backlog due to split discovery throttling to avoid
+ // excessive memory footprint. Some pending splits may not have been discovered yet.
+ .setUnassignedSplitsGauge(() -> Long.valueOf(splitProvider.pendingSplitCount()));
+ }
+
+ @Override
+ public void start() {
+ splitProvider.start();
+ }
+
+ @Override
+ public void handleSplitRequest(int subtaskId, @Nullable String requesterHostname) {
+ readersAwaitingSplit.put(subtaskId, requesterHostname);
+ assignSplits();
+ }
+
+ @Override
+ public void handleSourceEvent(int subtaskId, SourceEvent sourceEvent) {
+ if (sourceEvent instanceof SplitRequestEvent) {
+ SplitRequestEvent splitRequestEvent = (SplitRequestEvent) sourceEvent;
+ LOG.info("Received request split event from subtask {}", subtaskId);
+ splitProvider.onCompletedSplits(splitRequestEvent.finishedSplitIds());
+ readersAwaitingSplit.put(subtaskId, splitRequestEvent.requesterHostname());
+ assignSplits();
+ } else {
+ throw new IllegalArgumentException(
+ String.format(
+ Locale.ROOT,
+ "Received unknown event from subtask %d: %s",
+ subtaskId,
+ sourceEvent.getClass().getCanonicalName()));
+ }
+ }
+
+ @Override
+ public void addSplitsBack(List<HoodieSourceSplit> splits, int subtaskId) {
+ LOG.info("Add {} splits back to the split provider for failed subtask {}", splits.size(), subtaskId);
+ splitProvider.onUnassignedSplits(splits);
+ assignSplits();
+ }
+
+ @Override
+ public void addReader(int subtaskId) {
+ LOG.info("Add reader for subtask {}", subtaskId);
+ }
+
+ @Override
+ public HoodieSplitEnumeratorState snapshotState(long checkpointId) throws Exception {
+ return new HoodieSplitEnumeratorState(splitProvider.state());
+ }
+
+ @Override
+ public void close() throws IOException {
+ splitProvider.close();
+ }
+
+ @Override
+ public void handleSourceEvent(int subTaskId, int attemptNumber, SourceEvent sourceEvent) {
+ handleSourceEvent(subTaskId, sourceEvent);
+ }
+
+ private void assignSplits() {
+ LOG.info("Assigning splits for {} awaiting readers", readersAwaitingSplit.size());
+ Iterator<Map.Entry<Integer, String>> awaitingReader =
+ readersAwaitingSplit.entrySet().iterator();
+ while (awaitingReader.hasNext()) {
+ Map.Entry<Integer, String> nextAwaiting = awaitingReader.next();
+ // if the reader that requested another split has failed in the meantime, remove
+ // it from the list of waiting readers
+ if (!enumeratorContext.registeredReaders().containsKey(nextAwaiting.getKey())) {
+ awaitingReader.remove();
+ continue;
+ }
+
+ int awaitingSubtask = nextAwaiting.getKey();
+ String hostname = nextAwaiting.getValue();
+ Option<HoodieSourceSplit> splitOption = splitProvider.getNext(hostname);
+ if (splitOption.isPresent()) {
+ LOG.info("Assign split to subtask {}: {}", awaitingSubtask, splitOption.get());
+ enumeratorContext.assignSplit(splitOption.get(), awaitingSubtask);
+ awaitingReader.remove();
+ } else if (splitOption.isEmpty()) {
+ if (shouldWaitForMoreSplits()) {
+ // for continuous split provider callback
+ registerCallbackFromSplitProvider();
+ break;
+ } else {
+ // for static split provider
+ LOG.info("No more splits available for subtask {}", awaitingSubtask);
+ enumeratorContext.signalNoMoreSplits(awaitingSubtask);
+ awaitingReader.remove();
+ }
+ }
+ }
+ }
+
+ abstract boolean shouldWaitForMoreSplits();
+
+ private synchronized void registerCallbackFromSplitProvider() {
+ if (availableFuture.get() != null) {
+ return;
+ }
+
+ CompletableFuture<Void> future =
+ splitProvider
+ .isAvailable()
+ .thenAccept(
+ ignore ->
+ enumeratorContext.runInCoordinatorThread(
+ () -> {
+ LOG.debug("Executing callback of new splits from provider");
+ availableFuture.set(null);
+ assignSplits();
+ }));
+ availableFuture.set(future);
+ LOG.debug("Registered callback for future available splits");
+ }
+}
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/enumerator/HoodieSplitEnumeratorState.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/enumerator/HoodieSplitEnumeratorState.java
new file mode 100644
index 000000000000..de41d1d0e0f4
--- /dev/null
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/enumerator/HoodieSplitEnumeratorState.java
@@ -0,0 +1,39 @@
+/*
+ * 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.hudi.source.enumerator;
+
+import org.apache.hudi.source.split.HoodieSourceSplitState;
+
+import java.io.Serializable;
+import java.util.Collection;
+
+/**
+ * State of Hoodie split enumerator. Mainly include the states of pending splits of split provider.
+ */
+public class HoodieSplitEnumeratorState implements Serializable {
+ private final Collection<HoodieSourceSplitState> pendingSplitStates;
+
+ public HoodieSplitEnumeratorState(Collection<HoodieSourceSplitState> pendingSplitStates) {
+ this.pendingSplitStates = pendingSplitStates;
+ }
+
+ public Collection<HoodieSourceSplitState> getPendingSplitStates() {
+ return pendingSplitStates;
+ }
+}
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/enumerator/HoodieStaticSplitEnumerator.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/enumerator/HoodieStaticSplitEnumerator.java
new file mode 100644
index 000000000000..49c0c67e054b
--- /dev/null
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/enumerator/HoodieStaticSplitEnumerator.java
@@ -0,0 +1,40 @@
+/*
+ * 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.hudi.source.enumerator;
+
+import org.apache.hudi.source.split.HoodieSourceSplit;
+import org.apache.hudi.source.split.HoodieSplitProvider;
+
+import org.apache.flink.api.connector.source.SplitEnumeratorContext;
+
+/**
+ * Static Hoodie split enumerator that only handles with a bounded number of hudi commits.
+ */
+public class HoodieStaticSplitEnumerator extends AbstractHoodieSplitEnumerator {
+
+ public HoodieStaticSplitEnumerator(
+ SplitEnumeratorContext<HoodieSourceSplit> enumeratorContext, HoodieSplitProvider provider) {
+ super(enumeratorContext, provider);
+ }
+
+ @Override
+ protected boolean shouldWaitForMoreSplits() {
+ return false;
+ }
+}
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/split/DefaultHoodieSplitProvider.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/split/DefaultHoodieSplitProvider.java
new file mode 100644
index 000000000000..5c29ac292aa8
--- /dev/null
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/split/DefaultHoodieSplitProvider.java
@@ -0,0 +1,84 @@
+/*
+ * 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.hudi.source.split;
+
+import org.apache.hudi.common.util.Option;
+
+import javax.annotation.Nullable;
+import java.util.Collection;
+import java.util.Queue;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentLinkedDeque;
+import java.util.stream.Collectors;
+
+/**
+ * Default split provider caches pending splits in queue and return split sequentially.
+ */
+public class DefaultHoodieSplitProvider implements HoodieSplitProvider {
+ private Queue<HoodieSourceSplit> pendingSplits;
+
+ public DefaultHoodieSplitProvider() {
+ this.pendingSplits = new ConcurrentLinkedDeque<>();
+ }
+
+ @Override
+ public Option<HoodieSourceSplit> getNext(@Nullable String hostname) {
+ if (!pendingSplits.isEmpty()) {
+ return Option.of(pendingSplits.poll());
+ }
+
+ return Option.empty();
+ }
+
+ @Override
+ public void onDiscoveredSplits(Collection<HoodieSourceSplit> splits) {
+ addSplits(splits);
+ }
+
+ @Override
+ public void onUnassignedSplits(Collection<HoodieSourceSplit> splits) {
+ addSplits(splits);
+ }
+
+ private void addSplits(Collection<HoodieSourceSplit> splits) {
+ this.pendingSplits.addAll(splits);
+ }
+
+ @Override
+ public Collection<HoodieSourceSplitState> state() {
+ return pendingSplits.stream().map(
+ split -> new HoodieSourceSplitState(split, HoodieSourceSplitStatus.UNASSIGNED))
+ .collect(Collectors.toList());
+ }
+
+ @Override
+ public CompletableFuture<Void> isAvailable() {
+ return new CompletableFuture<>();
+ }
+
+ @Override
+ public int pendingSplitCount() {
+ return pendingSplits.size();
+ }
+
+ @Override
+ public long pendingRecords() {
+ throw new UnsupportedOperationException("Pending records is not supported in DefaultSplitProvider.");
+ }
+}
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/split/HoodieSplitProvider.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/split/HoodieSplitProvider.java
new file mode 100644
index 000000000000..e4fa424c6741
--- /dev/null
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/split/HoodieSplitProvider.java
@@ -0,0 +1,122 @@
+/*
+ * 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.hudi.source.split;
+
+import org.apache.hudi.common.util.Option;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.connector.source.SplitEnumeratorContext;
+
+import javax.annotation.Nullable;
+
+import java.io.Closeable;
+import java.util.Collection;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * SplitProvider interface is to define split assignment strategy for different requirements. E.g.
+ *
+ * <ul>
+ * <li>Simple assigner with no ordering guarantee.
+ * <li>Instant time aware assigner that assign splits based on the order they are committed.
+ * </ul>
+ *
+ * <p>SplitProvider implementation needs to be thread safe. Enumerator call the assigner APIs mostly from
+ * the coordinator thread. But enumerator may call the {@link HoodieSplitProvider#pendingSplitCount()}
+ * from the I/O threads.
+ */
+@Internal
+public interface HoodieSplitProvider extends Closeable {
+
+ /**
+ * Some SplitProviders may need to start background threads or perform other activity such as
+ * registering as listeners to updates from other event sources e.g., watermark tracker.
+ */
+ default void start() {
+ }
+
+ /**
+ * Some SplitProviders may need to perform certain actions when their corresponding enumerators are
+ * closed
+ */
+ @Override
+ default void close() {
+ }
+
+ /**
+ * Request a new split from the assigner when enumerator trying to assign splits to awaiting
+ * readers.
+ *
+ * <p>If enumerator wasn't able to assign the split (e.g., reader disconnected), enumerator should
+ * call {@link HoodieSplitProvider#onUnassignedSplits} to return the split.
+ */
+ Option<HoodieSourceSplit> getNext(@Nullable String hostname);
+
+ /** Add new splits discovered by enumerator */
+ void onDiscoveredSplits(Collection<HoodieSourceSplit> splits);
+
+ /** Forward addSplitsBack event (for failed reader) to assigner */
+ void onUnassignedSplits(Collection<HoodieSourceSplit> splits);
+
+ /**
+ * Some assigner (like event time alignment) may track in-progress splits to advance watermark upon
+ * completed splits
+ */
+ default void onCompletedSplits(Collection<String> completedSplitIds) {
+ }
+
+ /**
+ * Get assigner state for checkpointing. This is a super-set API that works for all currently
+ * imagined assigners.
+ */
+ Collection<HoodieSourceSplitState> state();
+
+ /**
+ * Enumerator can get a notification via CompletableFuture when the assigner has more splits
+ * available later. Enumerator should schedule assignment in the thenAccept action of the future.
+ *
+ * <p>Assigner will return the same future if this method is called again before the previous
+ * future is completed.
+ *
+ * <p>The future can be completed from other thread, e.g. the coordinator thread from another
+ * thread for event time alignment.
+ *
+ * <p>If enumerator need to trigger action upon the future completion, it may want to run it in
+ * the coordinator thread using {@link SplitEnumeratorContext#runInCoordinatorThread(Runnable)}.
+ */
+ CompletableFuture<Void> isAvailable();
+
+ /**
+ * Return the number of pending splits that haven't been assigned yet.
+ *
+ * <p>The enumerator can poll this API to publish a metric on the number of pending splits.
+ *
+ * <p>The enumerator can also use this information to throttle split discovery for streaming read.
+ * If there are already many pending splits tracked by the assigner, it is undesirable to discover
+ * more splits and track them in the assigner. That will increase the memory footprint and
+ * enumerator checkpoint size.
+ */
+ int pendingSplitCount();
+
+ /**
+ * Return the number of pending records, which can act as a measure of the source lag. This value
+ * could be an estimation if the exact number of records cannot be accurately computed.
+ */
+ long pendingRecords();
+}
diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/source/enumerator/TestHoodieStaticSplitEnumerator.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/source/enumerator/TestHoodieStaticSplitEnumerator.java
new file mode 100644
index 000000000000..cc640ac871a5
--- /dev/null
+++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/source/enumerator/TestHoodieStaticSplitEnumerator.java
@@ -0,0 +1,311 @@
+/*
+ * 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.hudi.source.enumerator;
+
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.source.split.DefaultHoodieSplitProvider;
+import org.apache.hudi.source.split.HoodieSourceSplit;
+import org.apache.hudi.source.split.SplitRequestEvent;
+
+import org.apache.flink.api.connector.source.ReaderInfo;
+import org.apache.flink.api.connector.source.SplitEnumeratorContext;
+import org.apache.flink.api.connector.source.SplitsAssignment;
+import org.apache.flink.metrics.groups.SplitEnumeratorMetricGroup;
+import org.apache.flink.metrics.groups.UnregisteredMetricsGroup;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Callable;
+import java.util.function.BiConsumer;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * Test cases for {@link HoodieStaticSplitEnumerator}.
+ */
+public class TestHoodieStaticSplitEnumerator {
+ private MockSplitEnumeratorContext context;
+ private DefaultHoodieSplitProvider splitProvider;
+ private HoodieStaticSplitEnumerator enumerator;
+ private HoodieSourceSplit split1;
+ private HoodieSourceSplit split2;
+ private HoodieSourceSplit split3;
+
+ @BeforeEach
+ public void setUp() {
+ context = new MockSplitEnumeratorContext();
+ splitProvider = new DefaultHoodieSplitProvider();
+ enumerator = new HoodieStaticSplitEnumerator(context, splitProvider);
+
+ split1 = createTestSplit(1, "file1");
+ split2 = createTestSplit(2, "file2");
+ split3 = createTestSplit(3, "file3");
+ }
+
+ @Test
+ public void testStartEnumerator() {
+ enumerator.start();
+ // Verify start doesn't throw exception
+ }
+
+ @Test
+ public void testHandleSplitRequest() {
+ splitProvider.onDiscoveredSplits(Arrays.asList(split1, split2));
+ enumerator.start();
+
+ context.registerReader(new ReaderInfo(0, "localhost"));
+ enumerator.handleSplitRequest(0, "localhost");
+
+ // Check that split was assigned
+ assertEquals(1, context.getAssignedSplits().size(), "Should have assigned splits");
+ assertTrue(context.getAssignedSplits().containsKey(0), "Should assign to subtask 0");
+ }
+
+ @Test
+ public void testHandleSourceEvent() {
+ splitProvider.onDiscoveredSplits(Arrays.asList(split1, split2));
+ enumerator.start();
+
+ context.registerReader(new ReaderInfo(0, "localhost"));
+
+ // Send split request event
+ SplitRequestEvent event = new SplitRequestEvent(Collections.emptyList(), "localhost");
+ enumerator.handleSourceEvent(0, event);
+
+ assertEquals(1, context.getAssignedSplits().size(), "Should assign split after receiving event");
+ }
+
+ @Test
+ public void testHandleSourceEventWithCompletedSplits() {
+ splitProvider.onDiscoveredSplits(Arrays.asList(split1, split2, split3));
+ enumerator.start();
+
+ context.registerReader(new ReaderInfo(0, "localhost"));
+
+ // First request
+ enumerator.handleSplitRequest(0, "localhost");
+ int initialCount = splitProvider.pendingSplitCount();
+
+ // Send event with completed splits
+ SplitRequestEvent event = new SplitRequestEvent(Arrays.asList(split1.splitId()), "localhost");
+ enumerator.handleSourceEvent(0, event);
+
+ // Should assign another split
+ assertTrue(context.getAssignedSplits().size() > 0, "Should have assigned splits");
+ }
+
+ @Test
+ public void testAddSplitsBack() {
+ splitProvider.onDiscoveredSplits(Arrays.asList(split1));
+ enumerator.start();
+
+ context.registerReader(new ReaderInfo(0, "localhost"));
+ enumerator.handleSplitRequest(0, "localhost");
+
+ // Simulate reader failure - add splits back
+ enumerator.addSplitsBack(Arrays.asList(split1), 0);
+
+ assertEquals(1, splitProvider.pendingSplitCount(), "Split should be added back to provider");
+ }
+
+ @Test
+ public void testAddReader() {
+ context.registerReader(new ReaderInfo(0, "localhost"));
+ enumerator.addReader(0);
+ // Should not throw exception
+ }
+
+ @Test
+ public void testSnapshotState() throws Exception {
+ splitProvider.onDiscoveredSplits(Arrays.asList(split1, split2));
+
+ HoodieSplitEnumeratorState state = enumerator.snapshotState(1L);
+
+ assertNotNull(state, "Snapshot state should not be null");
+ }
+
+ @Test
+ public void testNoMoreSplitsSignal() {
+ enumerator.start();
+ context.registerReader(new ReaderInfo(0, "localhost"));
+
+ // Request split when provider is empty
+ enumerator.handleSplitRequest(0, "localhost");
+
+ // Should signal no more splits for static enumerator
+ assertTrue(context.getNoMoreSplitsSignaled().contains(0),
+ "Should signal no more splits when no splits available");
+ }
+
+ @Test
+ public void testMultipleReaders() {
+ splitProvider.onDiscoveredSplits(Arrays.asList(split1, split2, split3));
+ enumerator.start();
+
+ // Register multiple readers
+ context.registerReader(new ReaderInfo(0, "host1"));
+ context.registerReader(new ReaderInfo(1, "host2"));
+
+ // Request splits for both readers
+ enumerator.handleSplitRequest(0, "host1");
+ enumerator.handleSplitRequest(1, "host2");
+
+ // Both readers should get splits
+ assertTrue(context.getAssignedSplits().containsKey(0), "Reader 0 should get a split");
+ assertTrue(context.getAssignedSplits().containsKey(1), "Reader 1 should get a split");
+ }
+
+ @Test
+ public void testReaderDisconnectedBeforeAssignment() {
+ splitProvider.onDiscoveredSplits(Arrays.asList(split1));
+ enumerator.start();
+
+ context.registerReader(new ReaderInfo(0, "localhost"));
+ enumerator.handleSplitRequest(0, "localhost");
+
+ // Unregister reader
+ context.unregisterReader(0);
+
+ // Request another split (shouldn't assign to disconnected reader)
+ splitProvider.onDiscoveredSplits(Arrays.asList(split2));
+ enumerator.handleSplitRequest(0, "localhost");
+
+ // Reader 0 should not receive the second split since it's disconnected
+ List<HoodieSourceSplit> assignedToReader0 = context.getAssignedSplits().get(0);
+ if (assignedToReader0 != null) {
+ assertEquals(1, assignedToReader0.size(), "Disconnected reader should not receive new splits");
+ }
+ }
+
+ @Test
+ public void testHandleSourceEventWithAttemptNumber() {
+ splitProvider.onDiscoveredSplits(Arrays.asList(split1));
+ enumerator.start();
+
+ context.registerReader(new ReaderInfo(0, "localhost"));
+
+ SplitRequestEvent event = new SplitRequestEvent(Collections.emptyList(), "localhost");
+ enumerator.handleSourceEvent(0, 1, event);
+
+ assertTrue(context.getAssignedSplits().size() > 0, "Should assign split via attempt-aware method");
+ }
+
+ private HoodieSourceSplit createTestSplit(int splitNum, String fileId) {
+ return new HoodieSourceSplit(
+ splitNum,
+ "basePath_" + splitNum,
+ Option.empty(),
+ "/table/path",
+ "read_optimized",
+ fileId
+ );
+ }
+
+ /**
+ * Mock implementation of SplitEnumeratorContext for testing.
+ */
+ private static class MockSplitEnumeratorContext implements SplitEnumeratorContext<HoodieSourceSplit> {
+ private final Map<Integer, ReaderInfo> registeredReaders = new HashMap<>();
+ private final Map<Integer, List<HoodieSourceSplit>> assignedSplits = new HashMap<>();
+ private final List<Integer> noMoreSplitsSignaled = new ArrayList<>();
+ private final List<Runnable> coordinatorThreadTasks = new ArrayList<>();
+
+ public void registerReader(ReaderInfo readerInfo) {
+ registeredReaders.put(readerInfo.getSubtaskId(), readerInfo);
+ }
+
+ public void unregisterReader(int subtaskId) {
+ registeredReaders.remove(subtaskId);
+ }
+
+ public Map<Integer, List<HoodieSourceSplit>> getAssignedSplits() {
+ return assignedSplits;
+ }
+
+ public List<Integer> getNoMoreSplitsSignaled() {
+ return noMoreSplitsSignaled;
+ }
+
+ @Override
+ public SplitEnumeratorMetricGroup metricGroup() {
+ return UnregisteredMetricsGroup.createSplitEnumeratorMetricGroup();
+ }
+
+ @Override
+ public void sendEventToSourceReader(int subtaskId, org.apache.flink.api.connector.source.SourceEvent event) {
+ // No-op for testing
+ }
+
+ @Override
+ public int currentParallelism() {
+ return registeredReaders.size();
+ }
+
+ @Override
+ public Map<Integer, ReaderInfo> registeredReaders() {
+ return new HashMap<>(registeredReaders);
+ }
+
+ @Override
+ public void assignSplits(SplitsAssignment<HoodieSourceSplit> newSplitAssignments) {
+ newSplitAssignments.assignment().forEach((subtask, splits) -> {
+ assignedSplits.computeIfAbsent(subtask, k -> new ArrayList<>()).addAll(splits);
+ });
+ }
+
+ @Override
+ public void assignSplit(HoodieSourceSplit split, int subtask) {
+ assignedSplits.computeIfAbsent(subtask, k -> new ArrayList<>()).add(split);
+ }
+
+ @Override
+ public void signalNoMoreSplits(int subtask) {
+ noMoreSplitsSignaled.add(subtask);
+ }
+
+ @Override
+ public <T> void callAsync(Callable<T> callable, BiConsumer<T, Throwable> handler) {
+ try {
+ T result = callable.call();
+ handler.accept(result, null);
+ } catch (Exception e) {
+ handler.accept(null, e);
+ }
+ }
+
+ @Override
+ public <T> void callAsync(Callable<T> callable, BiConsumer<T, Throwable> handler, long initialDelay, long period) {
+ callAsync(callable, handler);
+ }
+
+ @Override
+ public void runInCoordinatorThread(Runnable runnable) {
+ coordinatorThreadTasks.add(runnable);
+ runnable.run();
+ }
+ }
+}
diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/source/split/TestDefaultHoodieSplitProvider.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/source/split/TestDefaultHoodieSplitProvider.java
new file mode 100644
index 000000000000..8b413b23622d
--- /dev/null
+++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/source/split/TestDefaultHoodieSplitProvider.java
@@ -0,0 +1,188 @@
+/*
+ * 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.hudi.source.split;
+
+import org.apache.hudi.common.util.Option;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * Test cases for {@link DefaultHoodieSplitProvider}.
+ */
+public class TestDefaultHoodieSplitProvider {
+ private DefaultHoodieSplitProvider provider;
+ private HoodieSourceSplit split1;
+ private HoodieSourceSplit split2;
+ private HoodieSourceSplit split3;
+
+ @BeforeEach
+ public void setUp() {
+ provider = new DefaultHoodieSplitProvider();
+ split1 = createTestSplit(1, "file1");
+ split2 = createTestSplit(2, "file2");
+ split3 = createTestSplit(3, "file3");
+ }
+
+ @Test
+ public void testGetNextFromEmptyProvider() {
+ Option<HoodieSourceSplit> result = provider.getNext(null);
+ assertFalse(result.isPresent(), "Should return empty option when no splits available");
+ }
+
+ @Test
+ public void testGetNextWithHostname() {
+ provider.onDiscoveredSplits(Arrays.asList(split1, split2));
+
+ Option<HoodieSourceSplit> result = provider.getNext("localhost");
+ assertTrue(result.isPresent(), "Should return a split");
+ assertEquals(split1.splitId(), result.get().splitId(), "Should return first split");
+ }
+
+ @Test
+ public void testGetNextSequentially() {
+ provider.onDiscoveredSplits(Arrays.asList(split1, split2, split3));
+
+ Option<HoodieSourceSplit> first = provider.getNext(null);
+ assertTrue(first.isPresent(), "First split should be present");
+ assertEquals(split1.splitId(), first.get().splitId(), "Should return splits in order");
+
+ Option<HoodieSourceSplit> second = provider.getNext(null);
+ assertTrue(second.isPresent(), "Second split should be present");
+ assertEquals(split2.splitId(), second.get().splitId(), "Should return splits in order");
+
+ Option<HoodieSourceSplit> third = provider.getNext(null);
+ assertTrue(third.isPresent(), "Third split should be present");
+ assertEquals(split3.splitId(), third.get().splitId(), "Should return splits in order");
+
+ Option<HoodieSourceSplit> empty = provider.getNext(null);
+ assertFalse(empty.isPresent(), "Should return empty when all splits consumed");
+ }
+
+ @Test
+ public void testOnDiscoveredSplits() {
+ List<HoodieSourceSplit> splits = Arrays.asList(split1, split2);
+ provider.onDiscoveredSplits(splits);
+
+ assertEquals(2, provider.pendingSplitCount(), "Should have 2 pending splits");
+
+ provider.getNext(null);
+ assertEquals(1, provider.pendingSplitCount(), "Should have 1 pending split after consuming one");
+ }
+
+ @Test
+ public void testOnUnassignedSplits() {
+ provider.onDiscoveredSplits(Arrays.asList(split1));
+ provider.getNext(null);
+
+ assertEquals(0, provider.pendingSplitCount(), "Should have no pending splits");
+
+ // Re-add unassigned split
+ provider.onUnassignedSplits(Arrays.asList(split1));
+ assertEquals(1, provider.pendingSplitCount(), "Should have 1 pending split after re-adding");
+
+ Option<HoodieSourceSplit> result = provider.getNext(null);
+ assertTrue(result.isPresent(), "Should be able to get the re-added split");
+ assertEquals(split1.splitId(), result.get().splitId(), "Should return the re-added split");
+ }
+
+ @Test
+ public void testState() {
+ provider.onDiscoveredSplits(Arrays.asList(split1, split2, split3));
+
+ Collection<HoodieSourceSplitState> states = provider.state();
+ assertEquals(3, states.size(), "Should have 3 split states");
+
+ // Verify all states are UNASSIGNED
+ for (HoodieSourceSplitState state : states) {
+ assertEquals(HoodieSourceSplitStatus.UNASSIGNED, state.status(),
+ "All splits should have UNASSIGNED status");
+ }
+ }
+
+ @Test
+ public void testStateAfterConsumingSomeSplits() {
+ provider.onDiscoveredSplits(Arrays.asList(split1, split2, split3));
+
+ // Consume first two splits
+ provider.getNext(null);
+ provider.getNext(null);
+
+ Collection<HoodieSourceSplitState> states = provider.state();
+ assertEquals(1, states.size(), "Should have 1 remaining split state");
+
+ HoodieSourceSplitState state = states.iterator().next();
+ assertEquals(split3.splitId(), state.split().splitId(), "Remaining split should be split3");
+ assertEquals(HoodieSourceSplitStatus.UNASSIGNED, state.status(), "Split should be UNASSIGNED");
+ }
+
+ @Test
+ public void testPendingSplitCount() {
+ assertEquals(0, provider.pendingSplitCount(), "Initially should have 0 pending splits");
+
+ provider.onDiscoveredSplits(Arrays.asList(split1, split2));
+ assertEquals(2, provider.pendingSplitCount(), "Should have 2 pending splits");
+
+ provider.getNext(null);
+ assertEquals(1, provider.pendingSplitCount(), "Should have 1 pending split");
+
+ provider.getNext(null);
+ assertEquals(0, provider.pendingSplitCount(), "Should have 0 pending splits");
+ }
+
+ @Test
+ public void testIsAvailable() {
+ // isAvailable returns a non-completed future for DefaultSplitProvider
+ assertTrue(provider.isAvailable() != null, "isAvailable should return a future");
+ assertFalse(provider.isAvailable().isDone(), "Future should not be completed");
+ }
+
+ @Test
+ public void testMultipleDiscoveredSplitsCalls() {
+ provider.onDiscoveredSplits(Arrays.asList(split1));
+ assertEquals(1, provider.pendingSplitCount(), "Should have 1 split");
+
+ provider.onDiscoveredSplits(Arrays.asList(split2, split3));
+ assertEquals(3, provider.pendingSplitCount(), "Should have 3 splits total");
+
+ // Verify order is maintained
+ assertEquals(split1.splitId(), provider.getNext(null).get().splitId());
+ assertEquals(split2.splitId(), provider.getNext(null).get().splitId());
+ assertEquals(split3.splitId(), provider.getNext(null).get().splitId());
+ }
+
+ private HoodieSourceSplit createTestSplit(int splitNum, String fileId) {
+ return new HoodieSourceSplit(
+ splitNum,
+ "basePath_" + splitNum,
+ Option.empty(),
+ "/table/path",
+ "read_optimized",
+ fileId
+ );
+ }
+}