This is an automated email from the ASF dual-hosted git repository.
davidzollo pushed a commit to branch dev
in repository https://gitbox.apache.org/repos/asf/seatunnel.git
The following commit(s) were added to refs/heads/dev by this push:
new e2bb766710 [Improve][connector-starrocks] Improved starrocks source
enumerator splits allocation algorithm for subtasks (#10867)
e2bb766710 is described below
commit e2bb766710ddaca333c1aa011013488c7d7e1200
Author: JeremyXin <[email protected]>
AuthorDate: Thu May 28 12:03:42 2026 +0800
[Improve][connector-starrocks] Improved starrocks source enumerator splits
allocation algorithm for subtasks (#10867)
---
.../starrocks/source/StarRocksSourceState.java | 1 +
.../source/StartRocksSourceSplitEnumerator.java | 27 ++-
.../source/StarRocksSourceSplitEnumeratorTest.java | 209 +++++++++++++++++++++
3 files changed, 231 insertions(+), 6 deletions(-)
diff --git
a/seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/source/StarRocksSourceState.java
b/seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/source/StarRocksSourceState.java
index d19451cc18..cc96323097 100644
---
a/seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/source/StarRocksSourceState.java
+++
b/seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/source/StarRocksSourceState.java
@@ -33,4 +33,5 @@ public class StarRocksSourceState implements Serializable {
private static final long serialVersionUID = -147928488869915694L;
private Map<Integer, List<StarRocksSourceSplit>> pendingSplit;
private final ConcurrentLinkedQueue<String> pendingTables;
+ private int assignCount;
}
diff --git
a/seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/source/StartRocksSourceSplitEnumerator.java
b/seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/source/StartRocksSourceSplitEnumerator.java
index f3c0f855d9..62bd964dc4 100644
---
a/seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/source/StartRocksSourceSplitEnumerator.java
+++
b/seatunnel-connectors-v2/connector-starrocks/src/main/java/org/apache/seatunnel/connectors/seatunnel/starrocks/source/StartRocksSourceSplitEnumerator.java
@@ -29,11 +29,13 @@ import lombok.extern.slf4j.Slf4j;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
+import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
@Slf4j
@@ -46,6 +48,8 @@ public class StartRocksSourceSplitEnumerator
private final Object stateLock = new Object();
private final Context<StarRocksSourceSplit> context;
+ // Keeps round-robin assignment global across tables and restore.
+ private final AtomicInteger assignCount = new AtomicInteger(0);
public StartRocksSourceSplitEnumerator(
SourceSplitEnumerator.Context<StarRocksSourceSplit> context,
@@ -71,6 +75,7 @@ public class StartRocksSourceSplitEnumerator
if (sourceState != null) {
this.pendingSplit.putAll(sourceState.getPendingSplit());
this.pendingTables.addAll(sourceState.getPendingTables());
+ this.assignCount.set(sourceState.getAssignCount());
}
}
@@ -99,7 +104,7 @@ public class StartRocksSourceSplitEnumerator
public void addSplitsBack(List<StarRocksSourceSplit> splits, int
subtaskId) {
log.debug("Add back splits {} to StartRocksSourceSplitEnumerator.",
splits);
if (!splits.isEmpty()) {
- addPendingSplit(splits);
+ addPendingSplit(splits, subtaskId);
assignSplit(Collections.singletonList(subtaskId));
}
}
@@ -120,7 +125,7 @@ public class StartRocksSourceSplitEnumerator
@Override
public StarRocksSourceState snapshotState(long checkpointId) {
synchronized (stateLock) {
- return new StarRocksSourceState(pendingSplit, pendingTables);
+ return new StarRocksSourceState(pendingSplit, pendingTables,
assignCount.get());
}
}
@@ -147,13 +152,23 @@ public class StartRocksSourceSplitEnumerator
private void addPendingSplit(Collection<StarRocksSourceSplit> splits) {
int readerCount = context.currentParallelism();
- for (StarRocksSourceSplit split : splits) {
- int ownerReader = getSplitOwner(split.splitId(), readerCount);
+
+ List<StarRocksSourceSplit> sortedSplits =
+ splits.stream()
+
.sorted(Comparator.comparing(StarRocksSourceSplit::getSplitId))
+ .collect(Collectors.toList());
+
+ for (StarRocksSourceSplit split : sortedSplits) {
+ int ownerReader = getSplitOwner(assignCount.getAndIncrement(),
readerCount);
log.info("Assigning {} to {} reader.", split.getSplitId(),
ownerReader);
pendingSplit.computeIfAbsent(ownerReader, r -> new
ArrayList<>()).add(split);
}
}
+ private void addPendingSplit(Collection<StarRocksSourceSplit> splits, int
ownerReader) {
+ pendingSplit.computeIfAbsent(ownerReader, r -> new
ArrayList<>()).addAll(splits);
+ }
+
private void assignSplit(Collection<Integer> readers) {
log.debug("Assign pendingSplits to readers {}", readers);
@@ -191,7 +206,7 @@ public class StartRocksSourceSplitEnumerator
return sourceSplits;
}
- private static int getSplitOwner(String tp, int numReaders) {
- return (tp.hashCode() & Integer.MAX_VALUE) % numReaders;
+ private static int getSplitOwner(int assignCount, int numReaders) {
+ return assignCount % numReaders;
}
}
diff --git
a/seatunnel-connectors-v2/connector-starrocks/src/test/java/org/apache/seatunnel/connectors/seatunnel/starrocks/source/StarRocksSourceSplitEnumeratorTest.java
b/seatunnel-connectors-v2/connector-starrocks/src/test/java/org/apache/seatunnel/connectors/seatunnel/starrocks/source/StarRocksSourceSplitEnumeratorTest.java
new file mode 100644
index 0000000000..d17f45272d
--- /dev/null
+++
b/seatunnel-connectors-v2/connector-starrocks/src/test/java/org/apache/seatunnel/connectors/seatunnel/starrocks/source/StarRocksSourceSplitEnumeratorTest.java
@@ -0,0 +1,209 @@
+/*
+ * 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.seatunnel.connectors.seatunnel.starrocks.source;
+
+import org.apache.seatunnel.api.configuration.ReadonlyConfig;
+import org.apache.seatunnel.api.source.SourceSplitEnumerator;
+import org.apache.seatunnel.connectors.seatunnel.starrocks.config.SourceConfig;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentLinkedQueue;
+
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyList;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+
+public class StarRocksSourceSplitEnumeratorTest {
+
+ @Test
+ public void shouldBalanceSplitsEvenlyAcrossReaders() throws Exception {
+ Map<Integer, List<StarRocksSourceSplit>> assignments = new HashMap<>();
+ SourceSplitEnumerator.Context<StarRocksSourceSplit> context =
createContext(4, assignments);
+ StartRocksSourceSplitEnumerator enumerator =
+ new StartRocksSourceSplitEnumerator(
+ context, new
SourceConfig(ReadonlyConfig.fromMap(configMap())));
+
+ Method addPendingSplit =
+ StartRocksSourceSplitEnumerator.class.getDeclaredMethod(
+ "addPendingSplit", java.util.Collection.class);
+ addPendingSplit.setAccessible(true);
+ addPendingSplit.invoke(enumerator, buildSplits(10));
+
+ StarRocksSourceState state = enumerator.snapshotState(1L);
+ Map<Integer, List<StarRocksSourceSplit>> pendingSplits =
state.getPendingSplit();
+
+ Assertions.assertEquals(4, pendingSplits.size());
+ Assertions.assertEquals(3, pendingSplits.get(0).size());
+ Assertions.assertEquals(3, pendingSplits.get(1).size());
+ Assertions.assertEquals(2, pendingSplits.get(2).size());
+ Assertions.assertEquals(2, pendingSplits.get(3).size());
+ }
+
+ @Test
+ public void shouldBalanceSingleSplitTablesAcrossReadersInRun() {
+ Map<Integer, List<StarRocksSourceSplit>> assignments = new HashMap<>();
+ SourceSplitEnumerator.Context<StarRocksSourceSplit> context =
createContext(3, assignments);
+ StartRocksSourceSplitEnumerator enumerator =
+ spy(
+ new StartRocksSourceSplitEnumerator(
+ context,
+ new
SourceConfig(ReadonlyConfig.fromMap(configMap())),
+ new StarRocksSourceState(
+ new HashMap<>(),
+ new ConcurrentLinkedQueue<>(
+ new ArrayList<String>() {
+ {
+ add("table_0");
+ add("table_1");
+ add("table_2");
+ }
+ }),
+ 0)));
+
+ doReturn(Collections.singletonList(buildSplit("split-0")))
+ .when(enumerator)
+ .getStarRocksSourceSplit("table_0");
+ doReturn(Collections.singletonList(buildSplit("split-1")))
+ .when(enumerator)
+ .getStarRocksSourceSplit("table_1");
+ doReturn(Collections.singletonList(buildSplit("split-2")))
+ .when(enumerator)
+ .getStarRocksSourceSplit("table_2");
+
+ enumerator.run();
+
+ Assertions.assertEquals(1, getAssignedSplitCount(assignments, 0));
+ Assertions.assertEquals(1, getAssignedSplitCount(assignments, 1));
+ Assertions.assertEquals(1, getAssignedSplitCount(assignments, 2));
+ }
+
+ @Test
+ public void shouldReassignReturnedSplitsToOriginalReader() {
+ Map<Integer, List<StarRocksSourceSplit>> assignments = new HashMap<>();
+ SourceSplitEnumerator.Context<StarRocksSourceSplit> context =
createContext(3, assignments);
+ StartRocksSourceSplitEnumerator enumerator =
+ new StartRocksSourceSplitEnumerator(
+ context, new
SourceConfig(ReadonlyConfig.fromMap(configMap())));
+
+ enumerator.addSplitsBack(buildSplits(3), 2);
+
+ Assertions.assertEquals(0, getAssignedSplitCount(assignments, 0));
+ Assertions.assertEquals(0, getAssignedSplitCount(assignments, 1));
+ Assertions.assertEquals(3, getAssignedSplitCount(assignments, 2));
+
+ StarRocksSourceState state = enumerator.snapshotState(1L);
+ Assertions.assertTrue(state.getPendingSplit().isEmpty());
+ }
+
+ @Test
+ public void shouldContinueRoundRobinAfterRestore() {
+ Map<Integer, List<StarRocksSourceSplit>> assignments = new HashMap<>();
+ SourceSplitEnumerator.Context<StarRocksSourceSplit> context =
createContext(3, assignments);
+ String remainingTable = "table_after_restore";
+ StartRocksSourceSplitEnumerator enumerator =
+ spy(
+ new StartRocksSourceSplitEnumerator(
+ context,
+ new
SourceConfig(ReadonlyConfig.fromMap(configMap())),
+ new StarRocksSourceState(
+ new HashMap<>(),
+ new ConcurrentLinkedQueue<>(
+
Collections.singletonList(remainingTable)),
+ 1)));
+
+ doReturn(Collections.singletonList(buildSplit("split-after-restore")))
+ .when(enumerator)
+ .getStarRocksSourceSplit(remainingTable);
+
+ enumerator.run();
+
+ Assertions.assertEquals(0, getAssignedSplitCount(assignments, 0));
+ Assertions.assertEquals(1, getAssignedSplitCount(assignments, 1));
+ Assertions.assertEquals(0, getAssignedSplitCount(assignments, 2));
+ }
+
+ private Map<String, Object> configMap() {
+ Map<String, Object> config = new HashMap<>();
+ config.put("nodeUrls", Collections.singletonList("localhost:8030"));
+ config.put("username", "root");
+ config.put("password", "");
+ config.put("database", "test_db");
+ config.put("table", "test_table");
+ config.put("table_list", Collections.emptyList());
+ return config;
+ }
+
+ private List<StarRocksSourceSplit> buildSplits(int size) {
+ List<StarRocksSourceSplit> splits = new ArrayList<>();
+ for (int i = 0; i < size; i++) {
+ splits.add(buildSplit("split-" + i));
+ }
+ return splits;
+ }
+
+ private StarRocksSourceSplit buildSplit(String splitId) {
+ return new StarRocksSourceSplit(null, splitId);
+ }
+
+ private SourceSplitEnumerator.Context<StarRocksSourceSplit> createContext(
+ int parallelism, Map<Integer, List<StarRocksSourceSplit>>
assignments) {
+ @SuppressWarnings("unchecked")
+ SourceSplitEnumerator.Context<StarRocksSourceSplit> context =
+ mock(SourceSplitEnumerator.Context.class);
+ when(context.currentParallelism()).thenReturn(parallelism);
+
when(context.registeredReaders()).thenReturn(createReaders(parallelism));
+ doAnswer(
+ invocation -> {
+ int subtaskId = invocation.getArgument(0);
+ List<StarRocksSourceSplit> splits =
invocation.getArgument(1);
+ assignments
+ .computeIfAbsent(subtaskId, ignored -> new
ArrayList<>())
+ .addAll(splits);
+ return null;
+ })
+ .when(context)
+ .assignSplit(anyInt(), anyList());
+ return context;
+ }
+
+ private Set<Integer> createReaders(int parallelism) {
+ Set<Integer> readers = new java.util.LinkedHashSet<>();
+ for (int i = 0; i < parallelism; i++) {
+ readers.add(i);
+ }
+ return readers;
+ }
+
+ private int getAssignedSplitCount(
+ Map<Integer, List<StarRocksSourceSplit>> assignments, int
subtaskId) {
+ return assignments.getOrDefault(subtaskId,
Collections.emptyList()).size();
+ }
+}
