[HUDI-4924] Auto-tune dedup parallelism (#6802)

Author: yihua

hudi-client/hudi-client-common/src/main/java/org/apache/hudi/table/action/commit/HoodieWriteHelper.java

@@ -56,6 +56,9 @@ public HoodieData<HoodieRecord<T>> deduplicateRecords(
       HoodieData<HoodieRecord<T>> records, HoodieIndex<?, ?> index, int parallelism, String schemaStr) {
     boolean isIndexingGlobal = index.isGlobal();
     final SerializableSchema schema = new SerializableSchema(schemaStr);
+    // Auto-tunes the parallelism for reduce transformation based on the number of data partitions
+    // in engine-specific representation
+    int reduceParallelism = Math.max(1, Math.min(records.getNumPartitions(), parallelism));
     return records.mapToPair(record -> {
       HoodieKey hoodieKey = record.getKey();
       // If index used is global, then records are expected to differ in their partitionPath
@@ -67,7 +70,7 @@ public HoodieData<HoodieRecord<T>> deduplicateRecords(
       HoodieKey reducedKey = rec1.getData().equals(reducedData) ? rec1.getKey() : rec2.getKey();
 
       return new HoodieAvroRecord<>(reducedKey, reducedData);
-    }, parallelism).map(Pair::getRight);
+    }, reduceParallelism).map(Pair::getRight);
   }
 
 }

hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/data/HoodieJavaRDD.java

@@ -102,6 +102,11 @@ public long count() {
     return rddData.count();
   }
 
+  @Override
+  public int getNumPartitions() {
+    return rddData.getNumPartitions();
+  }
+
   @Override
   public <O> HoodieData<O> map(SerializableFunction<T, O> func) {
     return HoodieJavaRDD.of(rddData.map(func::apply));

hudi-client/hudi-spark-client/src/test/java/org/apache/hudi/client/functional/TestHoodieClientOnCopyOnWriteStorage.java

@@ -461,22 +461,27 @@ private void testDeduplication(
     HoodieData<HoodieRecord<RawTripTestPayload>> records = HoodieJavaRDD.of(
         jsc.parallelize(Arrays.asList(recordOne, recordTwo, recordThree), 1));
     HoodieWriteConfig.Builder configBuilder = getConfigBuilder(HoodieFailedWritesCleaningPolicy.LAZY)
-            .combineInput(true, true);
+        .combineInput(true, true);
     addConfigsForPopulateMetaFields(configBuilder, populateMetaFields);
     HoodieWriteConfig writeConfig = configBuilder.build();
 
     // Global dedup should be done based on recordKey only
     HoodieIndex index = mock(HoodieIndex.class);
     when(index.isGlobal()).thenReturn(true);
-    List<HoodieRecord<RawTripTestPayload>> dedupedRecs = HoodieWriteHelper.newInstance().deduplicateRecords(records, index, 1, writeConfig.getSchema()).collectAsList();
+    int dedupParallelism = records.getNumPartitions() + 100;
+    HoodieData<HoodieRecord<RawTripTestPayload>> dedupedRecsRdd = HoodieWriteHelper.newInstance().deduplicateRecords(records, index, dedupParallelism, writeConfig.getSchema());
+    List<HoodieRecord<RawTripTestPayload>> dedupedRecs = dedupedRecsRdd.collectAsList();
+    assertEquals(records.getNumPartitions(), dedupedRecsRdd.getNumPartitions());
     assertEquals(1, dedupedRecs.size());
     assertEquals(dedupedRecs.get(0).getPartitionPath(), recordThree.getPartitionPath());
     assertNodupesWithinPartition(dedupedRecs);
 
     // non-Global dedup should be done based on both recordKey and partitionPath
     index = mock(HoodieIndex.class);
     when(index.isGlobal()).thenReturn(false);
-    dedupedRecs = HoodieWriteHelper.newInstance().deduplicateRecords(records, index, 1, writeConfig.getSchema()).collectAsList();
+    dedupedRecsRdd = HoodieWriteHelper.newInstance().deduplicateRecords(records, index, dedupParallelism, writeConfig.getSchema());
+    dedupedRecs = dedupedRecsRdd.collectAsList();
+    assertEquals(records.getNumPartitions(), dedupedRecsRdd.getNumPartitions());
     assertEquals(2, dedupedRecs.size());
     assertNodupesWithinPartition(dedupedRecs);
 

hudi-client/hudi-spark-client/src/test/java/org/apache/hudi/data/TestHoodieJavaRDD.java

@@ -0,0 +1,40 @@
+/*
+ * 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.hudi.data;
+
+import org.apache.hudi.common.data.HoodieData;
+import org.apache.hudi.testutils.HoodieClientTestBase;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class TestHoodieJavaRDD extends HoodieClientTestBase {
+  @Test
+  public void testGetNumPartitions() {
+    int numPartitions = 6;
+    HoodieData<Integer> rddData = HoodieJavaRDD.of(jsc.parallelize(
+        IntStream.rangeClosed(0, 100).boxed().collect(Collectors.toList()), numPartitions));
+    assertEquals(numPartitions, rddData.getNumPartitions());
+  }
+}

hudi-common/src/main/java/org/apache/hudi/common/data/HoodieData.java

@@ -67,14 +67,19 @@ public interface HoodieData<T> extends Serializable {
 
   /**
    * Returns number of objects held in the collection
-   *
+   * <p>
    * NOTE: This is a terminal operation
    */
   long count();
 
+  /**
+   * @return the number of data partitions in the engine-specific representation.
+   */
+  int getNumPartitions();
+
   /**
    * Maps every element in the collection using provided mapping {@code func}.
-   *
+   * <p>
    * This is an intermediate operation
    *
    * @param func serializable map function

hudi-common/src/main/java/org/apache/hudi/common/data/HoodieListData.java

@@ -175,6 +175,11 @@ public long count() {
     return super.count();
   }
 
+  @Override
+  public int getNumPartitions() {
+    return 1;
+  }
+
   @Override
   public List<T> collectAsList() {
     return super.collectAsList();

hudi-common/src/test/java/org/apache/hudi/common/data/TestHoodieListData.java

@@ -29,6 +29,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -64,4 +65,11 @@ void testEagerSemantic() {
     assertEquals(3, originalListData.count());
     assertEquals(sourceList, originalListData.collectAsList());
   }
+
+  @Test
+  public void testGetNumPartitions() {
+    HoodieData<Integer> listData = HoodieListData.eager(
+        IntStream.rangeClosed(0, 100).boxed().collect(Collectors.toList()));
+    assertEquals(1, listData.getNumPartitions());
+  }
 }