[Enhancement] Get partitions by batch of size 5k per RPC (#19241)

(cherry picked from commit 7401b49)
StarRocks · Mar 10, 2023 · 69a59b7 · 69a59b7
1 parent f9c9434
commit 69a59b7
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/common/Config.java b/fe/fe-core/src/main/java/com/starrocks/common/Config.java
@@ -1610,7 +1610,7 @@ public class Config extends ConfigBase {
      * The maximum number of partitions to fetch from the metastore in one RPC.
      */
     @ConfField
-    public static int max_hive_partitions_per_rpc = 1000;
+    public static int max_hive_partitions_per_rpc = 5000;
 
     /**
      * The interval of lazy refreshing remote file's metadata cache

diff --git a/fe/fe-core/src/main/java/com/starrocks/connector/hive/HiveMetastore.java b/fe/fe-core/src/main/java/com/starrocks/connector/hive/HiveMetastore.java
@@ -30,6 +30,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
@@ -101,8 +102,15 @@ public Partition getPartition(String dbName, String tblName, List<String> partit
     }
 
     public Map<String, Partition> getPartitionsByNames(String dbName, String tblName, List<String> partitionNames) {
-        List<org.apache.hadoop.hive.metastore.api.Partition> partitions =
-                client.getPartitionsByNames(dbName, tblName, partitionNames);
+        List<org.apache.hadoop.hive.metastore.api.Partition> partitions = new ArrayList<>();
+        // fetch partitions by batch per RPC
+        for (int start = 0; start < partitionNames.size(); start += Config.max_hive_partitions_per_rpc) {
+            int end = Math.min(start + Config.max_hive_partitions_per_rpc, partitionNames.size());
+            List<String> namesPerRPC = partitionNames.subList(start, end);
+            List<org.apache.hadoop.hive.metastore.api.Partition> partsPerRPC =
+                    client.getPartitionsByNames(dbName, tblName, namesPerRPC);
+            partitions.addAll(partsPerRPC);
+        }
 
         Map<String, List<String>> partitionNameToPartitionValues = partitionNames.stream()
                 .collect(Collectors.toMap(Function.identity(), PartitionUtil::toPartitionValues));