[Improve][connector-doris] Improved doris source enumerator splits allocation algorithm for subtasks

pom.xml

@@ -568,6 +568,13 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-inline</artifactId>
+            <version>${mockito.version}</version>
+            <scope>test</scope>
+        </dependency>
+
         <!-- The prometheus simpleclient -->
         <dependency>
             <groupId>io.prometheus</groupId>

seatunnel-connectors-v2/connector-doris/src/main/java/org/apache/seatunnel/connectors/doris/rest/PartitionDefinition.java

@@ -129,7 +129,7 @@ public int hashCode() {
     @Override
     public String toString() {
         return "PartitionDefinition{"
-                + ", database='"
+                + "database='"
                 + database
                 + '\''
                 + ", table='"

seatunnel-connectors-v2/connector-doris/src/main/java/org/apache/seatunnel/connectors/doris/source/split/DorisSourceSplitEnumerator.java

@@ -33,10 +33,13 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 
 @Slf4j
 public class DorisSourceSplitEnumerator
@@ -52,6 +55,8 @@ public class DorisSourceSplitEnumerator
     private final Map<TablePath, DorisSourceTable> dorisSourceTables;
     private final Object stateLock = new Object();
 
+    private final AtomicInteger assignCount = new AtomicInteger(0);
+
     public DorisSourceSplitEnumerator(
             Context<DorisSourceSplit> context,
             DorisSourceConfig dorisSourceConfig,
@@ -162,15 +167,24 @@ private List<DorisSourceSplit> getDorisSourceSplit() {
 
     private void addPendingSplit(Collection<DorisSourceSplit> splits) {
         int readerCount = context.currentParallelism();
-        for (DorisSourceSplit split : splits) {
-            int ownerReader = getSplitOwner(split.splitId(), readerCount);
+
+        // sorting the splits to ensure the order
+        List<DorisSourceSplit> sortedSplits =
+                splits.stream()
+                        .sorted(Comparator.comparing(DorisSourceSplit::getSplitId))
+                        .collect(Collectors.toList());
+
+        // allocate splits in load balancing mode
+        assignCount.set(0);
+        for (DorisSourceSplit split : sortedSplits) {
+            int ownerReader = getSplitOwner(assignCount.getAndIncrement(), readerCount);
             log.info("Assigning split {} to reader {} .", split.splitId(), ownerReader);
             pendingSplit.computeIfAbsent(ownerReader, f -> new ArrayList<>()).add(split);
         }
     }
 
-    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;
     }
 
     private void assignSplit(Collection<Integer> readers) {

seatunnel-connectors-v2/connector-doris/src/test/java/org/apache/seatunnel/connectors/doris/split/DorisSourceSplitEnumeratorTest.java

@@ -0,0 +1,144 @@
+/*
+ * 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.doris.split;
+
+import org.apache.seatunnel.shade.com.google.common.collect.Maps;
+
+import org.apache.seatunnel.api.source.SourceSplitEnumerator;
+import org.apache.seatunnel.api.table.catalog.TablePath;
+import org.apache.seatunnel.connectors.doris.config.DorisSourceConfig;
+import org.apache.seatunnel.connectors.doris.rest.PartitionDefinition;
+import org.apache.seatunnel.connectors.doris.rest.RestService;
+import org.apache.seatunnel.connectors.doris.source.DorisSourceTable;
+import org.apache.seatunnel.connectors.doris.source.split.DorisSourceSplit;
+import org.apache.seatunnel.connectors.doris.source.split.DorisSourceSplitEnumerator;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.mockito.ArgumentMatchers.any;
+
+@Slf4j
+public class DorisSourceSplitEnumeratorTest {
+
+    private static final String DATABASE = "default";
+    private static final String TABLE = "default_table";
+    private static final String BE_ADDRESS_PREFIX = "doris-be-";
+    private static final String QUERY_PLAN = "DAABDAACDwABDAAAAAEIAA";
+
+    private static final int PARALLELISM = 4;
+
+    private static final int PARTITION_NUMS = 10;
+
+    @Test
+    public void dorisSourceSplitEnumeratorTest() {
+        DorisSourceConfig dorisSourceConfig = Mockito.mock(DorisSourceConfig.class);
+        DorisSourceTable dorisSourceTable = Mockito.mock(DorisSourceTable.class);
+
+        SourceSplitEnumerator.Context<DorisSourceSplit> context =
+                Mockito.mock(SourceSplitEnumerator.Context.class);
+
+        Mockito.when(context.registeredReaders())
+                .thenReturn(IntStream.range(0, PARALLELISM).boxed().collect(Collectors.toSet()));
+        Mockito.when(context.currentParallelism()).thenReturn(PARALLELISM);
+
+        Map<TablePath, DorisSourceTable> dorisSourceTableMap = Maps.newHashMap();
+        dorisSourceTableMap.put(new TablePath(DATABASE, null, TABLE), dorisSourceTable);
+
+        DorisSourceSplitEnumerator dorisSourceSplitEnumerator =
+                new DorisSourceSplitEnumerator(context, dorisSourceConfig, dorisSourceTableMap);
+
+        MockedStatic<RestService> restServiceMockedStatic = Mockito.mockStatic(RestService.class);
+
+        restServiceMockedStatic
+                .when(() -> RestService.findPartitions(any(), any(), any()))
+                .thenReturn(buildPartitionDefinitions());
+
+        dorisSourceSplitEnumerator.run();
+
+        ArgumentCaptor<Integer> subtaskId = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<List> split = ArgumentCaptor.forClass(List.class);
+
+        Mockito.verify(context, Mockito.times(PARALLELISM))
+                .assignSplit(subtaskId.capture(), split.capture());
+
+        List<Integer> subTaskAllValues = subtaskId.getAllValues();
+        List<List> splitAllValues = split.getAllValues();
+
+        for (int i = 0; i < PARALLELISM; i++) {
+            Assertions.assertEquals(i, subTaskAllValues.get(i));
+            Assertions.assertEquals(
+                    allocateFiles(i, PARALLELISM, PARTITION_NUMS), splitAllValues.get(i).size());
+        }
+
+        // check no duplicate file assigned
+        Assertions.assertEquals(0, dorisSourceSplitEnumerator.currentUnassignedSplitSize());
+    }
+
+    private List<PartitionDefinition> buildPartitionDefinitions() {
+
+        List<PartitionDefinition> partitions = new ArrayList<>();
+
+        IntStream.range(0, PARTITION_NUMS)
+                .forEach(
+                        i -> {
+                            PartitionDefinition partitionDefinition =
+                                    new PartitionDefinition(
+                                            DATABASE,
+                                            TABLE,
+                                            BE_ADDRESS_PREFIX + i,
+                                            new HashSet<>(i),
+                                            QUERY_PLAN);
+
+                            partitions.add(partitionDefinition);
+                        });
+
+        return partitions;
+    }
+
+    /**
+     * calculate the number of files assigned each time
+     *
+     * @param id id
+     * @param parallelism parallelism
+     * @param fileSize file size
+     * @return
+     */
+    public int allocateFiles(int id, int parallelism, int fileSize) {
+        int filesPerIteration = fileSize / parallelism;
+        int remainder = fileSize % parallelism;
+
+        if (id < remainder) {
+            return filesPerIteration + 1;
+        } else {
+            return filesPerIteration;
+        }
+    }
+}

seatunnel-translation/seatunnel-translation-base/pom.xml

@@ -37,5 +37,12 @@
             <version>${project.version}</version>
             <scope>test</scope>
         </dependency>
+
+        <dependency>
+            <groupId>org.apache.seatunnel</groupId>
+            <artifactId>connector-doris</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 </project>