nsivabalan commented on a change in pull request #1868: URL: https://github.com/apache/hudi/pull/1868#discussion_r463553042
########## File path: hudi-client/src/main/java/org/apache/hudi/table/action/commit/InsertBucket2CumulativeWeight.java ########## @@ -0,0 +1,53 @@ +/* + * 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.table.action.commit; + +import org.apache.hudi.common.util.collection.Pair; + +public class InsertBucket2CumulativeWeight extends Pair<InsertBucket, Double> { Review comment: may be InsertBucketCumulativeWeightPair ########## File path: hudi-client/src/main/java/org/apache/hudi/table/action/commit/UpsertPartitioner.java ########## @@ -193,15 +194,17 @@ private void assignInserts(WorkloadProfile profile, JavaSparkContext jsc) { } // Go over all such buckets, and assign weights as per amount of incoming inserts. - List<InsertBucket> insertBuckets = new ArrayList<>(); + List<Pair<InsertBucket, Double>> insertBuckets = new ArrayList<>(); Review comment: we might as well store InsertBucket2CumulativeWeight in this list since you have a class. ########## File path: hudi-client/src/main/java/org/apache/hudi/table/action/commit/UpsertPartitioner.java ########## @@ -270,20 +273,25 @@ public int getPartition(Object key) { return updateLocationToBucket.get(location.getFileId()); } else { String partitionPath = keyLocation._1().getPartitionPath(); - List<InsertBucket> targetBuckets = partitionPathToInsertBuckets.get(partitionPath); + List<Pair<InsertBucket, Double>> targetBuckets = partitionPathToInsertBucketInfos.get(partitionPath); // pick the target bucket to use based on the weights. - double totalWeight = 0.0; final long totalInserts = Math.max(1, profile.getWorkloadStat(partitionPath).getNumInserts()); final long hashOfKey = NumericUtils.getMessageDigestHash("MD5", keyLocation._1().getRecordKey()); final double r = 1.0 * Math.floorMod(hashOfKey, totalInserts) / totalInserts; - for (InsertBucket insertBucket : targetBuckets) { - totalWeight += insertBucket.weight; - if (r <= totalWeight) { - return insertBucket.bucketNumber; - } + + Pair<InsertBucket, Double> pair = new InsertBucket2CumulativeWeight(new InsertBucket(), r); Review comment: you don't need to declare. We could directly call. int index = Collections.binarySearch(targetBuckets, new InsertBucket2CumulativeWeight(new InsertBucket(), r)); ########## File path: hudi-client/src/main/java/org/apache/hudi/table/action/commit/UpsertPartitioner.java ########## @@ -272,21 +275,44 @@ public int getPartition(Object key) { String partitionPath = keyLocation._1().getPartitionPath(); List<InsertBucket> targetBuckets = partitionPathToInsertBuckets.get(partitionPath); // pick the target bucket to use based on the weights. - double totalWeight = 0.0; final long totalInserts = Math.max(1, profile.getWorkloadStat(partitionPath).getNumInserts()); final long hashOfKey = NumericUtils.getMessageDigestHash("MD5", keyLocation._1().getRecordKey()); final double r = 1.0 * Math.floorMod(hashOfKey, totalInserts) / totalInserts; - for (InsertBucket insertBucket : targetBuckets) { - totalWeight += insertBucket.weight; - if (r <= totalWeight) { - return insertBucket.bucketNumber; - } + + int index = binarySearch(targetBuckets, r); + if (index >= 0) { + return targetBuckets.get(index).bucketNumber; + } + + if (-1 * index - 1 < targetBuckets.size()) { Review comment: gotcha. can we add braces. ((-1*index) -1) in both if condition and in the return statement ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
