async and concurrently to process hudi table partitions in batches and create splits

plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/BridgingHiveMetastore.java

@@ -348,12 +348,13 @@ public Map<String, Optional<Partition>> getPartitionsByNames(Table table, List<S
         Map<List<String>, Partition> partitionValuesToPartitionMap = delegate.getPartitionsByNames(table.getDatabaseName(), table.getTableName(), partitionNames).stream()
                 .map(partition -> fromMetastoreApiPartition(table, partition))
                 .collect(Collectors.toMap(Partition::getValues, identity()));
-        ImmutableMap.Builder<String, Optional<Partition>> resultBuilder = ImmutableMap.builder();
+
+        ImmutableMap.Builder<String, Optional<Partition>> partitionNameToPartitionMap = ImmutableMap.builder();
         for (Map.Entry<String, List<String>> entry : partitionNameToPartitionValuesMap.entrySet()) {
             Partition partition = partitionValuesToPartitionMap.get(entry.getValue());
-            resultBuilder.put(entry.getKey(), Optional.ofNullable(partition));
+            partitionNameToPartitionMap.put(entry.getKey(), Optional.ofNullable(partition));
         }
-        return resultBuilder.buildOrThrow();
+        return partitionNameToPartitionMap.buildOrThrow();
     }
 
     private static Partition fromMetastoreApiPartition(Table table, org.apache.hadoop.hive.metastore.api.Partition partition)

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/ForHudiPartitionLoader.java

@@ -0,0 +1,34 @@
+/*
+ * Licensed 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 io.trino.plugin.hudi;
+
+import javax.inject.Qualifier;
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.FIELD;
+import static java.lang.annotation.ElementType.PARAMETER;
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+/**
+ * @author yaoxiao
+ * @version 1.0
+ * @date 2023/7/17 14:26
+ */
+
+@Retention(RUNTIME)
+@Target({FIELD, PARAMETER, METHOD})
+@Qualifier
+public @interface ForHudiPartitionLoader {}

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/ForHudiSplitLoader.java

@@ -0,0 +1,34 @@
+/*
+ * Licensed 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 io.trino.plugin.hudi;
+
+import javax.inject.Qualifier;
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.FIELD;
+import static java.lang.annotation.ElementType.PARAMETER;
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+/**
+ * @author yaoxiao
+ * @version 1.0
+ * @date 2023/7/17 14:26
+ */
+
+@Retention(RUNTIME)
+@Target({FIELD, PARAMETER, METHOD})
+@Qualifier
+public @interface ForHudiSplitLoader {}

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiConfig.java

@@ -39,13 +39,16 @@ public class HudiConfig
     private List<String> columnsToHide = ImmutableList.of();
     private boolean metadataEnabled;
     private boolean shouldUseParquetColumnNames = true;
-    private int minPartitionBatchSize = 10;
-    private int maxPartitionBatchSize = 100;
+    private int minPartitionBatchSize = 100;
+    private int maxPartitionBatchSize = 500;
     private boolean sizeBasedSplitWeightsEnabled = true;
     private DataSize standardSplitWeightSize = DataSize.of(128, MEGABYTE);
     private double minimumAssignedSplitWeight = 0.05;
     private int maxSplitsPerSecond = Integer.MAX_VALUE;
     private int maxOutstandingSplits = 1000;
+    private int partitionLoaderParallelism = 5;
+    private int splitLoaderParallelism = 5;
+
 
     public List<String> getColumnsToHide()
     {
@@ -191,4 +194,32 @@ public HudiConfig setMaxOutstandingSplits(int maxOutstandingSplits)
         this.maxOutstandingSplits = maxOutstandingSplits;
         return this;
     }
+
+    @Min(1)
+    public int getPartitionLoaderParallelism()
+    {
+        return partitionLoaderParallelism;
+    }
+
+    @Config("hudi.partition-loader-parallelism")
+    @ConfigDescription("hudi partition loader parallelism.")
+    public HudiConfig setPartitionLoaderParallelism(int partitionLoaderParallism)
+    {
+        this.partitionLoaderParallelism = partitionLoaderParallism;
+        return this;
+    }
+
+    @Min(1)
+    public int getSplitLoaderParallelism()
+    {
+        return splitLoaderParallelism;
+    }
+
+    @Config("hudi.split-loader-parallelism")
+    @ConfigDescription("hudi split loader parallelism.")
+    public HudiConfig setSplitLoaderParallelism(int splitLoaderParallelism)
+    {
+        this.splitLoaderParallelism = splitLoaderParallelism;
+        return this;
+    }
 }

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiErrorCode.java

@@ -28,7 +28,8 @@ public enum HudiErrorCode
     HUDI_MISSING_DATA(3, EXTERNAL),
     HUDI_CANNOT_OPEN_SPLIT(4, EXTERNAL),
     HUDI_UNSUPPORTED_FILE_FORMAT(5, EXTERNAL),
-    HUDI_CURSOR_ERROR(6, EXTERNAL);
+    HUDI_CURSOR_ERROR(6, EXTERNAL),
+    HUDI_FILESYSTEM_ERROR(7, EXTERNAL);
 
     private final ErrorCode errorCode;
 

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiMetadata.java

@@ -37,10 +37,11 @@
 import io.trino.spi.connector.TableNotFoundException;
 import io.trino.spi.predicate.TupleDomain;
 import io.trino.spi.type.TypeManager;
-import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.hudi.common.model.HoodieTableType;
 
+import java.io.IOException;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Iterator;
@@ -56,6 +57,7 @@
 import static io.trino.plugin.hive.util.HiveUtil.columnMetadataGetter;
 import static io.trino.plugin.hive.util.HiveUtil.hiveColumnHandles;
 import static io.trino.plugin.hive.util.HiveUtil.isHiveSystemSchema;
+import static io.trino.plugin.hudi.HudiErrorCode.HUDI_FILESYSTEM_ERROR;
 import static io.trino.plugin.hudi.HudiErrorCode.HUDI_UNKNOWN_TABLE_TYPE;
 import static io.trino.plugin.hudi.HudiSessionProperties.getColumnsToHide;
 import static io.trino.plugin.hudi.HudiTableProperties.LOCATION_PROPERTY;
@@ -65,7 +67,6 @@
 import static java.util.Collections.singletonList;
 import static java.util.Objects.requireNonNull;
 import static java.util.function.Function.identity;
-import static org.apache.hudi.common.fs.FSUtils.getFs;
 import static org.apache.hudi.common.table.HoodieTableMetaClient.METAFOLDER_NAME;
 import static org.apache.hudi.common.util.StringUtils.isNullOrEmpty;
 import static org.apache.hudi.exception.TableNotFoundException.checkTableValidity;
@@ -197,9 +198,15 @@ HiveMetastore getMetastore()
     private boolean isHudiTable(ConnectorSession session, Table table)
     {
         String basePath = table.getStorage().getLocation();
-        Configuration configuration = hdfsEnvironment.getConfiguration(new HdfsContext(session), new Path(basePath));
+        FileSystem fileSystem;
         try {
-            checkTableValidity(getFs(basePath, configuration), new Path(basePath), new Path(basePath, METAFOLDER_NAME));
+            fileSystem = hdfsEnvironment.getFileSystem(new HdfsContext(session), new Path(basePath));
+        } catch (IOException e) {
+            throw new TrinoException(HUDI_FILESYSTEM_ERROR, "Failed getting FileSystem: " + basePath, e);
+        }
+
+        try {
+            checkTableValidity(fileSystem, new Path(basePath), new Path(basePath, METAFOLDER_NAME));
         }
         catch (org.apache.hudi.exception.TableNotFoundException e) {
             log.warn("Could not find Hudi table at path '%s'", basePath);

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiSessionProperties.java

@@ -49,6 +49,8 @@ public class HudiSessionProperties
     private static final String SIZE_BASED_SPLIT_WEIGHTS_ENABLED = "size_based_split_weights_enabled";
     private static final String STANDARD_SPLIT_WEIGHT_SIZE = "standard_split_weight_size";
     private static final String MINIMUM_ASSIGNED_SPLIT_WEIGHT = "minimum_assigned_split_weight";
+    private static final String PARTITION_INFO_LOADER_PARALLELISM = "partition_info_loader_parallelism";
+    private static final String SPLIT_LOADER_PARALLELISM = "split_loader_parallelism";
 
     private final List<PropertyMetadata<?>> sessionProperties;
 
@@ -111,6 +113,16 @@ public HudiSessionProperties(HudiConfig hudiConfig, ParquetReaderConfig parquetR
                                 throw new TrinoException(INVALID_SESSION_PROPERTY, format("%s must be > 0 and <= 1.0: %s", MINIMUM_ASSIGNED_SPLIT_WEIGHT, value));
                             }
                         },
+                        false),
+                integerProperty(
+                        PARTITION_INFO_LOADER_PARALLELISM,
+                        "hudi partition info generator parallelism.",
+                        hudiConfig.getPartitionLoaderParallelism(),
+                        false),
+                integerProperty(
+                        SPLIT_LOADER_PARALLELISM,
+                        "hudi split generator parallelism.",
+                        hudiConfig.getSplitLoaderParallelism(),
                         false));
     }
 
@@ -165,4 +177,14 @@ public static double getMinimumAssignedSplitWeight(ConnectorSession session)
     {
         return session.getProperty(MINIMUM_ASSIGNED_SPLIT_WEIGHT, Double.class);
     }
+
+    public static int getPartitionInfoLoaderParallelism(ConnectorSession session)
+    {
+        return session.getProperty(PARTITION_INFO_LOADER_PARALLELISM, Integer.class);
+    }
+
+    public static int getSplitLoaderParallelism(ConnectorSession session)
+    {
+        return session.getProperty(SPLIT_LOADER_PARALLELISM, Integer.class);
+    }
 }

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiSplitManager.java

@@ -101,7 +101,8 @@ public ConnectorSplitSource getSplits(
                 partitionColumnHandles,
                 executor,
                 maxSplitsPerSecond,
-                maxOutstandingSplits);
+                maxOutstandingSplits,
+                hdfsEnvironment);
         return new ClassLoaderSafeConnectorSplitSource(splitSource, HudiSplitManager.class.getClassLoader());
     }
 }