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 63a7b529c5 [Feature][Connector-V2] Maxcompute Source Round-Robin Split
Assignment (#11131)
63a7b529c5 is described below
commit 63a7b529c588de68d88cfdfd6ee148480f772d02
Author: zhiliang-wu <[email protected]>
AuthorDate: Tue Jun 23 15:27:24 2026 +0200
[Feature][Connector-V2] Maxcompute Source Round-Robin Split Assignment
(#11131)
Co-authored-by: WU Zhiliang (External) <[email protected]>
---
.../source/MaxcomputeSourceSplitEnumerator.java | 61 +++++----
.../source/MaxcomputeSourceSplitTest.java | 151 +++++++++++++++++----
2 files changed, 162 insertions(+), 50 deletions(-)
diff --git
a/seatunnel-connectors-v2/connector-maxcompute/src/main/java/org/apache/seatunnel/connectors/seatunnel/maxcompute/source/MaxcomputeSourceSplitEnumerator.java
b/seatunnel-connectors-v2/connector-maxcompute/src/main/java/org/apache/seatunnel/connectors/seatunnel/maxcompute/source/MaxcomputeSourceSplitEnumerator.java
index 0f582bdd37..4dbd5abbd5 100644
---
a/seatunnel-connectors-v2/connector-maxcompute/src/main/java/org/apache/seatunnel/connectors/seatunnel/maxcompute/source/MaxcomputeSourceSplitEnumerator.java
+++
b/seatunnel-connectors-v2/connector-maxcompute/src/main/java/org/apache/seatunnel/connectors/seatunnel/maxcompute/source/MaxcomputeSourceSplitEnumerator.java
@@ -32,6 +32,7 @@ import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
+import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -108,37 +109,51 @@ public class MaxcomputeSourceSplitEnumerator
@Override
public void handleSplitRequest(int subtaskId) {}
+ // visible for testing
+ static Set<MaxcomputeSourceSplit> computeSplits(
+ int numReaders,
+ Collection<SourceTableInfo> sourceTableInfos,
+ Map<TablePath, Long> tableRecordCounts) {
+ Set<MaxcomputeSourceSplit> allSplit = new LinkedHashSet<>();
+ int chunkIndex = 0;
+ for (SourceTableInfo sourceTableInfo : sourceTableInfos) {
+ TablePath tablePath =
sourceTableInfo.getCatalogTable().getTablePath();
+ long recordCount = tableRecordCounts.get(tablePath);
+ int splitRow = MaxcomputeSourceOptions.SPLIT_ROW.defaultValue();
+ if (sourceTableInfo.getSplitRow() != null &&
sourceTableInfo.getSplitRow() > 0) {
+ splitRow = sourceTableInfo.getSplitRow();
+ }
+ for (long num = 0; num < recordCount; num += splitRow) {
+ int ownerReader = chunkIndex % numReaders;
+ allSplit.add(
+ new MaxcomputeSourceSplit(
+ num,
+ Math.min((long) splitRow, recordCount - num),
+ tablePath,
+ ownerReader));
+ chunkIndex++;
+ }
+ }
+ return allSplit;
+ }
+
private void discoverySplits() throws TunnelException {
int numReaders = enumeratorContext.currentParallelism();
- Set<MaxcomputeSourceSplit> allSplit = new HashSet<>();
+ Map<TablePath, Long> tableRecordCounts = new HashMap<>();
for (SourceTableInfo sourceTableInfo : sourceTableInfos.values()) {
- Set<MaxcomputeSourceSplit> splits = new HashSet<>();
TableTunnel.DownloadSession session =
MaxcomputeUtil.getDownloadSession(
readonlyConfig,
sourceTableInfo.getCatalogTable().getTablePath(),
sourceTableInfo.getPartitionSpec());
- long recordCount = session.getRecordCount();
- int splitRowNum = (int) Math.ceil((double) recordCount /
numReaders);
- int splitRow = MaxcomputeSourceOptions.SPLIT_ROW.defaultValue();
- if (sourceTableInfo.getSplitRow() != null &&
sourceTableInfo.getSplitRow() > 0) {
- splitRow = sourceTableInfo.getSplitRow();
- }
- for (int i = 0; i < numReaders; i++) {
- int readerStart = i * splitRowNum;
- int readerEnd = (int) Math.min((i + 1) * splitRowNum,
recordCount);
- for (int num = readerStart; num < readerEnd; num += splitRow) {
- splits.add(
- new MaxcomputeSourceSplit(
- num,
- Math.min(splitRow, readerEnd - num),
-
sourceTableInfo.getCatalogTable().getTablePath(),
- i));
- }
- }
- assignedSplits.forEach(splits::remove);
- allSplit.addAll(splits);
+ tableRecordCounts.put(
+ sourceTableInfo.getCatalogTable().getTablePath(),
session.getRecordCount());
}
+
+ Set<MaxcomputeSourceSplit> allSplit =
+ computeSplits(numReaders, sourceTableInfos.values(),
tableRecordCounts);
+ assignedSplits.forEach(allSplit::remove);
+
addSplitChangeToPendingAssignments(allSplit);
log.debug("Assigned {} to {} readers.", allSplit, numReaders);
log.info("Calculated splits successfully, the size of splits is {}.",
allSplit.size());
@@ -147,7 +162,7 @@ public class MaxcomputeSourceSplitEnumerator
private void
addSplitChangeToPendingAssignments(Collection<MaxcomputeSourceSplit> newSplits)
{
for (MaxcomputeSourceSplit split : newSplits) {
int ownerReader = split.getIndex() %
enumeratorContext.currentParallelism();
- pendingSplits.computeIfAbsent(ownerReader, r -> new
HashSet<>()).add(split);
+ pendingSplits.computeIfAbsent(ownerReader, r -> new
LinkedHashSet<>()).add(split);
}
}
diff --git
a/seatunnel-connectors-v2/connector-maxcompute/src/test/java/org/apache/seatunnel/connectors/seatunnel/maxcompute/source/MaxcomputeSourceSplitTest.java
b/seatunnel-connectors-v2/connector-maxcompute/src/test/java/org/apache/seatunnel/connectors/seatunnel/maxcompute/source/MaxcomputeSourceSplitTest.java
index d6dfcf6eba..83aae67404 100644
---
a/seatunnel-connectors-v2/connector-maxcompute/src/test/java/org/apache/seatunnel/connectors/seatunnel/maxcompute/source/MaxcomputeSourceSplitTest.java
+++
b/seatunnel-connectors-v2/connector-maxcompute/src/test/java/org/apache/seatunnel/connectors/seatunnel/maxcompute/source/MaxcomputeSourceSplitTest.java
@@ -17,14 +17,22 @@
package org.apache.seatunnel.connectors.seatunnel.maxcompute.source;
+import org.apache.seatunnel.api.table.catalog.CatalogTable;
import org.apache.seatunnel.api.table.catalog.TablePath;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
+import java.util.ArrayList;
+import java.util.HashMap;
import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
import java.util.Set;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
public class MaxcomputeSourceSplitTest {
@Test
@@ -33,42 +41,131 @@ public class MaxcomputeSourceSplitTest {
long recordCount = 105000;
int splitRow = 10000;
- int splitRowNum = (int) Math.ceil((double) recordCount / numReaders);
// 21000
TablePath tablePath = TablePath.of("test_db", "test_schema",
"test_table");
+ Map<TablePath, Long> tableRecordCounts = new HashMap<>();
+ tableRecordCounts.put(tablePath, recordCount);
+
+ CatalogTable catalogTable = mock(CatalogTable.class);
+ when(catalogTable.getTablePath()).thenReturn(tablePath);
+
+ SourceTableInfo tableInfo = new SourceTableInfo(catalogTable, null,
splitRow);
+ List<SourceTableInfo> tableInfos = new ArrayList<>();
+ tableInfos.add(tableInfo);
+
+ Set<MaxcomputeSourceSplit> splits =
+ MaxcomputeSourceSplitEnumerator.computeSplits(
+ numReaders, tableInfos, tableRecordCounts);
+
Set<String> splitIds = new HashSet<>();
int[] indexCounts = new int[numReaders];
- int totalSplits = 0;
-
- // Simulate the inner logic of
MaxcomputeSourceSplitEnumerator.discoverySplits
- for (int i = 0; i < numReaders; i++) {
- int readerStart = i * splitRowNum;
- int readerEnd = (int) Math.min((i + 1) * splitRowNum, recordCount);
- for (int num = readerStart; num < readerEnd; num += splitRow) {
- MaxcomputeSourceSplit split =
- new MaxcomputeSourceSplit(
- num, Math.min(splitRow, readerEnd - num),
tablePath, i);
-
- // Verify splitId uniqueness
- Assertions.assertTrue(
- splitIds.add(split.splitId()),
- "Duplicate splitId found: " + split.splitId());
- // Verify index assigned to the split matches the reader index
- Assertions.assertEquals(i, split.getIndex());
+ for (MaxcomputeSourceSplit split : splits) {
+ Assertions.assertTrue(
+ splitIds.add(split.splitId()), "Duplicate splitId found: "
+ split.splitId());
+ indexCounts[split.getIndex()]++;
+ }
- indexCounts[i]++;
- totalSplits++;
- }
+ Assertions.assertEquals(11, splits.size());
+
+ // Verify index distribution (round-robin):
+ // 11 splits / 5 readers = 2 splits each, plus 1 remainder for reader
0.
+ Assertions.assertEquals(3, indexCounts[0], "Reader 0 split count
mismatch!");
+ Assertions.assertEquals(2, indexCounts[1], "Reader 1 split count
mismatch!");
+ Assertions.assertEquals(2, indexCounts[2], "Reader 2 split count
mismatch!");
+ Assertions.assertEquals(2, indexCounts[3], "Reader 3 split count
mismatch!");
+ Assertions.assertEquals(2, indexCounts[4], "Reader 4 split count
mismatch!");
+ }
+
+ @Test
+ public void testMultiTableRoundRobinDistribution() {
+ int numReaders = 3;
+ int numTables = 5;
+ long recordCountPerTable = 15000;
+ int splitRow = 10000;
+
+ Map<TablePath, Long> tableRecordCounts = new HashMap<>();
+ List<SourceTableInfo> tableInfos = new ArrayList<>();
+
+ for (int t = 0; t < numTables; t++) {
+ TablePath tablePath = TablePath.of("test_db", "test_schema",
"test_table_" + t);
+ tableRecordCounts.put(tablePath, recordCountPerTable);
+
+ CatalogTable catalogTable = mock(CatalogTable.class);
+ when(catalogTable.getTablePath()).thenReturn(tablePath);
+
+ SourceTableInfo tableInfo = new SourceTableInfo(catalogTable,
null, splitRow);
+ tableInfos.add(tableInfo);
}
- Assertions.assertTrue(totalSplits == 15);
+ Set<MaxcomputeSourceSplit> splits =
+ MaxcomputeSourceSplitEnumerator.computeSplits(
+ numReaders, tableInfos, tableRecordCounts);
+
+ int[] indexCounts = new int[numReaders];
- // Verify index distribution:
- // Record count 105000 / 5 readers = 21000 splits per reader.
- // Split row is 10000. Loop steps: 0, 10000, 20000. So 3 splits per
reader.
- for (int i = 0; i < numReaders; i++) {
- Assertions.assertEquals(3, indexCounts[i], "Reader " + i + " split
count mismatch!");
+ for (MaxcomputeSourceSplit split : splits) {
+ indexCounts[split.getIndex()]++;
}
+
+ Assertions.assertEquals(10, splits.size());
+
+ // 10 splits / 3 readers = 3 splits each, plus 1 remainder for reader
0.
+ Assertions.assertEquals(4, indexCounts[0], "Reader 0 split count
mismatch!");
+ Assertions.assertEquals(3, indexCounts[1], "Reader 1 split count
mismatch!");
+ Assertions.assertEquals(3, indexCounts[2], "Reader 2 split count
mismatch!");
+ }
+
+ @Test
+ public void testComputeSplitsAscendingRowStartOrder() {
+ int numReaders = 3;
+ int numTables = 2;
+ Map<TablePath, Long> tableRecordCounts = new HashMap<>();
+ List<SourceTableInfo> tableInfos = new ArrayList<>();
+
+ for (int i = 0; i < numTables; i++) {
+ TablePath path = TablePath.of("db", "schema", "table" + i);
+ // 105,000 rows / 10,000 splitRow = 11 splits per table.
+ // With 3 readers, each reader gets 3-4 splits PER TABLE.
+ tableRecordCounts.put(path, 105000L);
+
+ CatalogTable catalogTable = mock(CatalogTable.class);
+ when(catalogTable.getTablePath()).thenReturn(path);
+
+ SourceTableInfo tableInfo = new SourceTableInfo(catalogTable,
null, 10000);
+ tableInfos.add(tableInfo);
+ }
+
+ // Drive the actual production logic
+ Set<MaxcomputeSourceSplit> splits =
+ MaxcomputeSourceSplitEnumerator.computeSplits(
+ numReaders, tableInfos, tableRecordCounts);
+
+ // Guard against HashSet regression: splits iterated from the Set MUST
be in insertion
+ // order,
+ // which means for any given reader and table, the rowStart must be
strictly ascending.
+ Map<Integer, Map<TablePath, Long>> lastRowStarts = new HashMap<>();
+
+ for (MaxcomputeSourceSplit split : splits) {
+ int reader = split.getIndex();
+ TablePath path = split.getTablePath();
+ long rowStart = split.getRowStart();
+
+ lastRowStarts.computeIfAbsent(reader, r -> new HashMap<>());
+ Long lastStart = lastRowStarts.get(reader).get(path);
+
+ if (lastStart != null) {
+ Assertions.assertTrue(
+ rowStart > lastStart,
+ "HashSet Regression! Splits for reader "
+ + reader
+ + " and table "
+ + path
+ + " are not in ascending order!");
+ }
+ lastRowStarts.get(reader).put(path, rowStart);
+ }
+
+ Assertions.assertEquals(22, splits.size(), "Should have exactly 22
splits total");
}
}
