(seatunnel) branch dev updated: [fix][connector-mango] fix split with avgSize zero error (#9255)

[wanghailin]

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

wanghailin 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 564863b933 [fix][connector-mango] fix split with avgSize zero error 
(#9255)
564863b933 is described below

commit 564863b933b0d9dc4d7492ec1b8c79e8c264cd2f
Author: CosmosNi <nijiahui_y...@cmss.chinamobile.com>
AuthorDate: Wed Apr 30 10:58:36 2025 +0800

    [fix][connector-mango] fix split with avgSize zero error (#9255)
---
 .../source/split/SamplingSplitStrategy.java        | 13 ++++++++
 .../source/split/SamplingSplitStrategyTest.java    | 38 ++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git 
a/seatunnel-connectors-v2/connector-mongodb/src/main/java/org/apache/seatunnel/connectors/seatunnel/mongodb/source/split/SamplingSplitStrategy.java
 
b/seatunnel-connectors-v2/connector-mongodb/src/main/java/org/apache/seatunnel/connectors/seatunnel/mongodb/source/split/SamplingSplitStrategy.java
index be37d38eda..8f320a6c42 100644
--- 
a/seatunnel-connectors-v2/connector-mongodb/src/main/java/org/apache/seatunnel/connectors/seatunnel/mongodb/source/split/SamplingSplitStrategy.java
+++ 
b/seatunnel-connectors-v2/connector-mongodb/src/main/java/org/apache/seatunnel/connectors/seatunnel/mongodb/source/split/SamplingSplitStrategy.java
@@ -76,6 +76,19 @@ public class SamplingSplitStrategy implements 
MongoSplitStrategy, Serializable {
         long count = numAndAvgSize.getLeft();
         long avgSize = numAndAvgSize.getRight();
 
+        // Handle the case when avgSize is 0 to prevent division by zero
+        if (avgSize <= 0) {
+            // If there are documents in the collection, return a single split
+            if (count > 0) {
+                return Lists.newArrayList(
+                        MongoSplitUtils.createMongoSplit(
+                                0, matchQuery, projection, splitKey, null, 
null));
+            } else {
+                // If there are no documents, return an empty list
+                return Lists.newArrayList();
+            }
+        }
+
         long numDocumentsPerSplit = sizePerSplit / avgSize;
         int numSplits = (int) Math.ceil((double) count / numDocumentsPerSplit);
         int numSamples = (int) Math.floor(samplesPerSplit * numSplits);
diff --git 
a/seatunnel-connectors-v2/connector-mongodb/src/test/java/org/apache/seatunnel/connectors/seatunnel/mongodb/source/split/SamplingSplitStrategyTest.java
 
b/seatunnel-connectors-v2/connector-mongodb/src/test/java/org/apache/seatunnel/connectors/seatunnel/mongodb/source/split/SamplingSplitStrategyTest.java
index ec5b05844f..926ed3803a 100644
--- 
a/seatunnel-connectors-v2/connector-mongodb/src/test/java/org/apache/seatunnel/connectors/seatunnel/mongodb/source/split/SamplingSplitStrategyTest.java
+++ 
b/seatunnel-connectors-v2/connector-mongodb/src/test/java/org/apache/seatunnel/connectors/seatunnel/mongodb/source/split/SamplingSplitStrategyTest.java
@@ -71,4 +71,42 @@ public class SamplingSplitStrategyTest {
         assertEquals(Long.valueOf(13360484963L), result.getLeft());
         assertEquals(Long.valueOf(200), result.getRight());
     }
+
+    @Test
+    public void testSplitWithZeroAvgSize() {
+        // Mock the getDocumentNumAndAvgSize method to return zero avgSize
+        SamplingSplitStrategy spyStrategy =
+                new SamplingSplitStrategy(
+                        clientProvider, "_id", new BsonDocument(), new 
BsonDocument(), 10L, 1024L) {
+                    @Override
+                    protected ImmutablePair<Long, Long> 
getDocumentNumAndAvgSize() {
+                        return ImmutablePair.of(10L, 0L); // 10 documents with 
0 avgSize
+                    }
+                };
+
+        // This should not throw a division by zero exception
+        java.util.List<MongoSplit> splits = spyStrategy.split();
+
+        // Should return a single split when count > 0 and avgSize = 0
+        assertEquals(1, splits.size());
+    }
+
+    @Test
+    public void testSplitWithZeroAvgSizeAndZeroCount() {
+        // Mock the getDocumentNumAndAvgSize method to return zero avgSize and 
zero count
+        SamplingSplitStrategy spyStrategy =
+                new SamplingSplitStrategy(
+                        clientProvider, "_id", new BsonDocument(), new 
BsonDocument(), 10L, 1024L) {
+                    @Override
+                    protected ImmutablePair<Long, Long> 
getDocumentNumAndAvgSize() {
+                        return ImmutablePair.of(0L, 0L); // 0 documents with 0 
avgSize
+                    }
+                };
+
+        // This should not throw a division by zero exception
+        java.util.List<MongoSplit> splits = spyStrategy.split();
+
+        // Should return an empty list when count = 0 and avgSize = 0
+        assertEquals(0, splits.size());
+    }
 }