Batched ingest files (#175)

snowflakedb · Jun 22, 2020 · 81f3554 · 81f3554
1 parent f3bdd10
commit 81f3554
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 10 deletions.
diff --git a/src/main/java/com/snowflake/kafka/connector/internal/SnowflakeIngestionService.java b/src/main/java/com/snowflake/kafka/connector/internal/SnowflakeIngestionService.java
@@ -27,9 +27,9 @@ public interface SnowflakeIngestionService
   /**
    * Ingest a list of files
    *
-   * @param fileNames file name set
+   * @param fileNames file name List
    */
-  void ingestFiles(Set<String> fileNames);
+  void ingestFiles(List<String> fileNames);
 
   /**
    * @return corresponding stage name

diff --git a/src/main/java/com/snowflake/kafka/connector/internal/SnowflakeIngestionServiceV1.java b/src/main/java/com/snowflake/kafka/connector/internal/SnowflakeIngestionServiceV1.java
@@ -6,11 +6,7 @@
 import net.snowflake.ingest.utils.StagedFileWrapper;
 
 import java.security.PrivateKey;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 
 import static com.snowflake.kafka.connector.internal.InternalUtils.convertIngestStatus;
 import static com.snowflake.kafka.connector.internal.InternalUtils.timestampToDate;
@@ -71,7 +67,7 @@ public void ingestFile(final String fileName)
   }
 
   @Override
-  public void ingestFiles(final Set<String> fileNames)
+  public void ingestFiles(final List<String> fileNames)
   {
     if (fileNames.isEmpty())
     {
@@ -80,7 +76,20 @@ public void ingestFiles(final Set<String> fileNames)
     try
     {
       InternalUtils.backoffAndRetry(telemetry,
-          () -> ingestManager.ingestFiles(SimpleIngestManager.wrapFilepaths(fileNames), null)
+          () ->
+          {
+            while (fileNames.size() > 0)
+            {
+              // Can not send more than 5000 files in one request,
+              // so batch 4000 as one request
+              int toIndex = Math.min(4000, fileNames.size());
+              List<String> fileNamesBatch = fileNames.subList(0, toIndex);
+              Set<String> fileNamesSet = new HashSet<>(fileNamesBatch);
+              ingestManager.ingestFiles(SimpleIngestManager.wrapFilepaths(fileNamesSet), null);
+              fileNamesBatch.clear();
+            }
+            return true;
+          }
       );
     } catch (Exception e)
     {

diff --git a/src/main/java/com/snowflake/kafka/connector/internal/SnowflakeSinkServiceV1.java b/src/main/java/com/snowflake/kafka/connector/internal/SnowflakeSinkServiceV1.java
@@ -490,7 +490,7 @@ private long getOffset()
         return committedOffset.get();
       }
 
-      Set<String> fileNamesCopy = new HashSet<>();
+      List<String> fileNamesCopy = new ArrayList<>();
       fileListLock.lock();
       try
       {